fix: argsz semantics in tests

Signed-off-by: William Henderson <william.henderson@nutanix.com>

2 files changed, 20 insertions, 16 deletions

diff --git a/test/py/test_dirty_pages.py b/test/py/test_dirty_pages.py
index a43352b..30fd757 100644
--- a/test/py/test_dirty_pages.py
+++ b/test/py/test_dirty_pages.py
@@ -144,7 +144,8 @@ def test_dirty_pages_start_different_pgsize():
 
 def test_dirty_pages_get_unmodified():
     argsz = len(vfio_user_device_feature()) + \
-            len(vfio_user_device_feature_dma_logging_report())
+            len(vfio_user_device_feature_dma_logging_report()) + \
+            8  # size of bitmap
 
     feature = vfio_user_device_feature(
         argsz=argsz,
@@ -161,11 +162,11 @@ def test_dirty_pages_get_unmodified():
 
     result = msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, payload)
 
-    assert len(result) == argsz + 8
+    assert len(result) == argsz
 
     feature, result = vfio_user_device_feature.pop_from_buffer(result)
 
-    assert feature.argsz == argsz + 8
+    assert feature.argsz == argsz
     assert feature.flags == VFIO_DEVICE_FEATURE_DMA_LOGGING_REPORT \
         | VFIO_DEVICE_FEATURE_GET
 
@@ -184,8 +185,11 @@ def test_dirty_pages_get_unmodified():
 
 def get_dirty_page_bitmap(addr=0x10 << PAGE_SHIFT, length=0x10 << PAGE_SHIFT,
                           page_size=PAGE_SIZE, expect=0):
+    bitmap_size = get_bitmap_size(length, page_size)
+
     argsz = len(vfio_user_device_feature()) + \
-            len(vfio_user_device_feature_dma_logging_report())
+            len(vfio_user_device_feature_dma_logging_report()) + \
+            bitmap_size
 
     feature = vfio_user_device_feature(
         argsz=argsz,
@@ -205,13 +209,13 @@ def get_dirty_page_bitmap(addr=0x10 << PAGE_SHIFT, length=0x10 << PAGE_SHIFT,
     if expect != 0:
         return
 
-    assert len(result) == argsz + get_bitmap_size(length, page_size)
+    assert len(result) == argsz
 
     _, result = vfio_user_device_feature.pop_from_buffer(result)
     _, result = \
         vfio_user_device_feature_dma_logging_report.pop_from_buffer(result)
 
-    assert len(result) == get_bitmap_size(length, page_size)
+    assert len(result) == bitmap_size
 
     return struct.unpack("Q", result)[0]
 
diff --git a/test/py/test_migration.py b/test/py/test_migration.py
index 0b36e1c..20de2bc 100644
--- a/test/py/test_migration.py
+++ b/test/py/test_migration.py
@@ -277,7 +277,7 @@ def test_handle_mig_data_read():
 
     transition_to_state(VFIO_USER_DEVICE_STATE_RUNNING)
 
-    argsz = len(vfio_user_mig_data())
+    argsz = len(vfio_user_mig_data()) + 4
 
     payload = vfio_user_mig_data(
         argsz=argsz,
@@ -289,14 +289,14 @@ def test_handle_mig_data_read():
     transition_to_state(VFIO_USER_DEVICE_STATE_PRE_COPY)
     read_data = data
     result = msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload)
-    assert len(result) == argsz + 4
-    assert result[argsz:argsz + 4] == data
+    assert len(result) == argsz
+    assert result[len(vfio_user_mig_data()):] == data
 
     transition_to_state(VFIO_USER_DEVICE_STATE_STOP_COPY)
     read_data = data
     result = msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload)
-    assert len(result) == argsz + 4
-    assert result[argsz:argsz + 4] == data
+    assert len(result) == argsz
+    assert result[len(vfio_user_mig_data()):] == data
 
 
 def test_handle_mig_data_read_too_long():
@@ -305,7 +305,7 @@ def test_handle_mig_data_read_too_long():
     transition_to_state(VFIO_USER_DEVICE_STATE_RUNNING)
 
     payload = vfio_user_mig_data(
-        argsz=len(vfio_user_mig_data()),
+        argsz=len(vfio_user_mig_data()) + 8,
         size=8
     )
 
@@ -322,7 +322,7 @@ def test_handle_mig_data_read_invalid_state():
     global read_data
 
     payload = vfio_user_mig_data(
-        argsz=len(vfio_user_mig_data()),
+        argsz=len(vfio_user_mig_data()) + 4,
         size=4
     )
 
@@ -345,7 +345,7 @@ def test_handle_mig_data_read_failed_callback():
     fail_callbacks = True
 
     payload = vfio_user_mig_data(
-        argsz=len(vfio_user_mig_data()),
+        argsz=len(vfio_user_mig_data()) + 4,
         size=4
     )
 
@@ -430,7 +430,7 @@ def test_handle_mig_data_write_short_write():
 
 def test_device_feature_migration():
     payload = vfio_user_device_feature(
-        argsz=len(vfio_user_device_feature()),
+        argsz=len(vfio_user_device_feature()) + 8,
         flags=VFIO_DEVICE_FEATURE_GET | VFIO_DEVICE_FEATURE_MIGRATION
     )
 
@@ -449,7 +449,7 @@ def test_device_feature_short_write():
 
 def test_device_feature_unsupported_operation():
     payload = vfio_user_device_feature(
-        argsz=len(vfio_user_device_feature()),
+        argsz=len(vfio_user_device_feature()) + 4,
         flags=VFIO_DEVICE_FEATURE_SET | VFIO_DEVICE_FEATURE_MIGRATION
     )