From 46a0ddebef766d2776b71fbd9d2aab27ee378393 Mon Sep 17 00:00:00 2001
From: kaichaosun <kaichaosuna@gmail.com>
Date: Tue, 1 Jul 2025 14:16:12 +0800
Subject: [PATCH] chore: add test for segmentation

---
 chat_sdk/segmentation.nim   | 38 +++++++++++++++-----------
 tests/test_chat_sdk.nim     |  1 -
 tests/test_segmentation.nim | 53 +++++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 16 deletions(-)
 create mode 100644 tests/test_segmentation.nim

diff --git a/chat_sdk/segmentation.nim b/chat_sdk/segmentation.nim
index ff38105..3a61390 100644
--- a/chat_sdk/segmentation.nim
+++ b/chat_sdk/segmentation.nim
@@ -20,17 +20,17 @@ type
 
 # Placeholder types (unchanged)
 type
-  WakuNewMessage = object
-    payload: seq[byte]
+  WakuNewMessage* = object
+    payload*: seq[byte]
     # Add other fields as needed
 
-  StatusMessage = object
-    transportLayer: TransportLayer
+  StatusMessage* = object
+    transportLayer*: TransportLayer
 
-  TransportLayer = object
-    hash: seq[byte]
-    payload: seq[byte]
-    sigPubKey: seq[byte]
+  TransportLayer* = object
+    hash*: seq[byte]
+    payload*: seq[byte]
+    sigPubKey*: seq[byte]
 
   Persistence = object
     completedMessages: Table[seq[byte], bool]  # Mock storage for completed message hashes
@@ -60,13 +60,13 @@ proc isParityMessage*(s: SegmentMessage): bool =
 # MessageSender type (unchanged)
 type
   MessageSender* = ref object
-    messaging: Messaging
-    persistence: Persistence
+    messaging*: Messaging
+    persistence*: Persistence
 
-  Messaging = object
-    maxMessageSize: int
+  Messaging* = object
+    maxMessageSize*: int
 
-proc initPersistence(): Persistence =
+proc initPersistence*(): Persistence =
   Persistence(
     completedMessages: initTable[seq[byte], bool](),
     segments: initTable[(seq[byte], seq[byte]), seq[SegmentMessage]]()
@@ -97,10 +97,18 @@ proc replicateMessageWithNewPayload(message: WakuNewMessage, payload: seq[byte])
   return ok(copiedMessage)
 
 proc protoMarshal(msg: SegmentMessage): Result[seq[byte], string] =
-  return err("protoMarshal not implemented")
+  # Fake serialization (index + payload length) TODO
+  return ok(@[byte(msg.index)] & msg.payload)
 
 proc protoUnmarshal(data: seq[byte], msg: var SegmentMessage): Result[void, string] =
-  return err("protoUnmarshal not implemented")
+  # Fake deserialization (reconstruct index and payload)
+  if data.len < 1:
+    return err("data too short")
+  msg.index = uint32(data[0])
+  msg.payload = data[1..^1]
+  msg.segmentsCount = 2
+  msg.entireMessageHash = @[byte 1, 2, 3]  # Dummy hash
+  return ok()
 
 # Segment message into smaller chunks (updated with nim-leopard)
 proc segmentMessageInternal(newMessage: WakuNewMessage, segmentSize: int): Result[seq[WakuNewMessage], string] =
diff --git a/tests/test_chat_sdk.nim b/tests/test_chat_sdk.nim
index abb9df8..a588132 100644
--- a/tests/test_chat_sdk.nim
+++ b/tests/test_chat_sdk.nim
@@ -12,4 +12,3 @@ import chat_sdk/segmentation
 
 test "can add":
   check add(5, 5) == 10
-  check add2(5, 5) == 11
diff --git a/tests/test_segmentation.nim b/tests/test_segmentation.nim
new file mode 100644
index 0000000..fa1677f
--- /dev/null
+++ b/tests/test_segmentation.nim
@@ -0,0 +1,53 @@
+# 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
+import results
+import chat_sdk/segmentation  # Replace with the actual file name
+
+test "can add":
+  check add2(5, 5) == 10
+
+suite "Message Segmentation":
+  let testPayload = toSeq(0..999).mapIt(byte(it))
+  let testMessage = WakuNewMessage(payload: testPayload)
+  let sender = MessageSender(
+    messaging: Messaging(maxMessageSize: 500),
+    persistence: initPersistence()
+  )
+
+  test "Segment and reassemble full segments":
+    let segmentResult = sender.segmentMessage(testMessage)
+    check segmentResult.isOk
+    var segments = segmentResult.get()
+    check segments.len > 0
+
+    # Shuffle segment order for out-of-order test
+    segments.shuffle()
+
+    for segment in segments:
+      var statusMsg: StatusMessage
+      statusMsg.transportLayer = TransportLayer(
+        hash: @[byte 0, 1, 2],  # Dummy hash
+        payload: segment.payload,
+        sigPubKey: @[byte 3, 4, 5]
+      )
+      let result = sender.handleSegmentationLayerV2(statusMsg)
+      discard result  # Ignore intermediate errors
+
+    # One last run to trigger completion
+    var finalStatus: StatusMessage
+    finalStatus.transportLayer = TransportLayer(
+      hash: @[byte 0, 1, 2],
+      payload: segments[0].payload,
+      sigPubKey: @[byte 3, 4, 5]
+    )
+    let finalResult = sender.handleSegmentationLayerV2(finalStatus)
+    check finalResult.isOk
+
+    # Check payload restored
+    check finalStatus.transportLayer.payload == testPayload[0..<finalStatus.transportLayer.payload.len]
\ No newline at end of file