nim-chat-sdk/tests/test_segmentation.nim

# This is just an example to get you started. You may wish to put all of your
# tests into a single file, or separate them into multiple `test1`, `test2`
# etc. files (better names are recommended, just make sure the name starts with
# the letter 't').
#
# To run these tests, simply execute `nimble test`.

import unittest, sequtils, random, math
import results
import db_connector/db_sqlite
import nimcrypto
import chat_sdk/segmentation
import chat_sdk/migration
import chat_sdk/db

proc newInMemoryPersistence(): SegmentationPersistence =
  let conn = open(":memory:", "", "", "")
  # Define the tables (same schema as expected in your app)
  runMigrations(conn)
  result = SegmentationPersistence(db: conn)

suite "message Segmentation":
  var
    sender: SegmentationHander
    testPayload: seq[byte]
    mockPersistence: SegmentationPersistence

  setup:
    # Initialize test payload (1000 bytes, each byte is its index)
    testPayload = newSeq[byte](1024)
    for i in 0..<1024:
      testPayload[i] = rand(255).byte

    # Setup mock persistence
    mockPersistence = newInMemoryPersistence()

    # Setup sender
    sender = SegmentationHander(
      segmentSize: 4000, # Arbitrary size to allow segmentation
      persistence: mockPersistence
    )
  
  test "HandleSegmentationLayer":
    # Define test cases (mirroring Go test cases)
    let testCases = @[
      (
        name: "all segments retrieved",
        segmentsCount: 2,
        expectedParitySegmentsCount: 0,
        retrievedSegments: @[0, 1],
        retrievedParitySegments: newSeq[int](),
        shouldSucceed: true
      ),
      (
        name: "all segments retrieved out of order",
        segmentsCount: 2,
        expectedParitySegmentsCount: 0,
        retrievedSegments: @[1, 0],
        retrievedParitySegments: newSeq[int](),
        shouldSucceed: true
      ),
      (
        name: "all segments&parity retrieved",
        segmentsCount: 8,
        expectedParitySegmentsCount: 1,
        retrievedSegments: @[0, 1, 2, 3, 4, 5, 6, 7, 8],
        retrievedParitySegments: @[8],
        shouldSucceed: true
      ),
      (
        name: "all segments&parity retrieved out of order",
        segmentsCount: 8,
        expectedParitySegmentsCount: 1,
        retrievedSegments: @[8, 0, 7, 1, 6, 2, 5, 3, 4],
        retrievedParitySegments: @[8],
        shouldSucceed: true
      ),
      (
        name: "no segments retrieved",
        segmentsCount: 2,
        expectedParitySegmentsCount: 0,
        retrievedSegments: @[],
        retrievedParitySegments: newSeq[int](),
        shouldSucceed: false
      ),
      (
        name: "not all needed segments&parity retrieved",
        segmentsCount: 8,
        expectedParitySegmentsCount: 1,
        retrievedSegments: @[1, 2, 8],
        retrievedParitySegments: @[8],
        shouldSucceed: false
      ),
      (
        name: "segments&parity retrieved",
        segmentsCount: 8,
        expectedParitySegmentsCount: 1,
        retrievedSegments: @[1, 2, 3, 4, 5, 6, 7, 8],
        retrievedParitySegments: @[8],
        shouldSucceed: true
      ),
      (
        name: "segments&parity retrieved out of order",
        segmentsCount: 16,
        expectedParitySegmentsCount: 2,
        retrievedSegments: @[17, 0, 16, 1, 15, 2, 14, 3, 13, 4, 12, 5, 11, 6, 10, 7],
        retrievedParitySegments: @[16, 17],
        shouldSucceed: true
      )
    ]

    for tc in testCases:
      test tc.name:
        # Segment the message
        let chunk = Chunk(payload: testPayload)
        sender.segmentSize = int(ceil(testPayload.len.float / tc.segmentsCount.float))
        let segmentedMessagesRes = sender.segmentMessage(chunk)
        require(segmentedMessagesRes.isOk)
        let segmentedMessages = segmentedMessagesRes.get()
        check(segmentedMessages.len == tc.segmentsCount + tc.expectedParitySegmentsCount)

        var message = Message(
          sigPubKey: keccak256.digest("testkey").data.toSeq,
          hash: keccak256.digest(testPayload).data.toSeq
        )

        var messageRecreated = false
        var handledSegments: seq[int] = @[]

        for i, segmentIndex in tc.retrievedSegments:
          message.payload = segmentedMessages[segmentIndex].payload
          let err = sender.handleSegmentationLayer(message)
          handledSegments.add(segmentIndex)
          if handledSegments.len < tc.segmentsCount:
            check(err.isErr and err.error == ErrMessageSegmentsIncomplete)
          elif handledSegments.len == tc.segmentsCount:
            check(err.isOk)
            check(message.payload == testPayload)
            messageRecreated = true
          else:
            check(err.isErr and err.error == ErrMessageSegmentsAlreadyCompleted)

        check(messageRecreated == tc.shouldSucceed)