test: improve test coverage

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

3 files changed, 176 insertions, 5 deletions

diff --git a/test/py/libvfio_user.py b/test/py/libvfio_user.py
index c83a75d..233f42a 100644
--- a/test/py/libvfio_user.py
+++ b/test/py/libvfio_user.py
@@ -211,11 +211,16 @@ VFIO_DEVICE_FEATURE_GET = (1 << 16)
 VFIO_DEVICE_FEATURE_SET = (1 << 17)
 VFIO_DEVICE_FEATURE_PROBE = (1 << 18)
 
+VFIO_DEVICE_FEATURE_MIGRATION = 1
 VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE = 2
 VFIO_DEVICE_FEATURE_DMA_LOGGING_START = 6
 VFIO_DEVICE_FEATURE_DMA_LOGGING_STOP = 7
 VFIO_DEVICE_FEATURE_DMA_LOGGING_REPORT = 8
 
+VFIO_MIGRATION_STOP_COPY = (1 << 0)
+VFIO_MIGRATION_P2P = (1 << 1)
+VFIO_MIGRATION_PRE_COPY = (1 << 2)
+
 VFIO_USER_IO_FD_TYPE_IOEVENTFD = 0
 VFIO_USER_IO_FD_TYPE_IOREGIONFD = 1
 VFIO_USER_IO_FD_TYPE_IOEVENTFD_SHADOW = 2
diff --git a/test/py/test_dirty_pages.py b/test/py/test_dirty_pages.py
index 3c13b8f..a43352b 100644
--- a/test/py/test_dirty_pages.py
+++ b/test/py/test_dirty_pages.py
@@ -102,12 +102,12 @@ def test_dirty_pages_setup():
     msg(ctx, sock, VFIO_USER_DMA_MAP, payload, fds=[f2.fileno()])
 
 
-def test_setup_migr_region():
+def test_setup_migration():
     ret = vfu_setup_device_migration_callbacks(ctx)
     assert ret == 0
 
 
-def start_logging(addr=None, length=None):
+def start_logging(addr=None, length=None, page_size=PAGE_SIZE, expect=0):
     if addr is not None:
         ranges = vfio_user_device_feature_dma_logging_range(
             iova=addr,
@@ -123,12 +123,12 @@ def start_logging(addr=None, length=None):
         flags=VFIO_DEVICE_FEATURE_DMA_LOGGING_START | VFIO_DEVICE_FEATURE_SET)
 
     payload = vfio_user_device_feature_dma_logging_control(
-        page_size=PAGE_SIZE,
+        page_size=page_size,
         num_ranges=(1 if addr is not None else 0),
         reserved=0)
 
     msg(ctx, sock, VFIO_USER_DEVICE_FEATURE,
-        bytes(feature) + bytes(payload) + bytes(ranges))
+        bytes(feature) + bytes(payload) + bytes(ranges), expect=expect)
 
 
 def test_dirty_pages_start():
@@ -137,6 +137,11 @@ def test_dirty_pages_start():
     start_logging()
 
 
+def test_dirty_pages_start_different_pgsize():
+    start_logging(page_size=0, expect=errno.EINVAL)
+    start_logging(page_size=PAGE_SIZE >> 1, expect=errno.EINVAL)
+
+
 def test_dirty_pages_get_unmodified():
     argsz = len(vfio_user_device_feature()) + \
             len(vfio_user_device_feature_dma_logging_report())
@@ -345,6 +350,22 @@ def test_dirty_pages_get_modified():
     assert bitmap == 0b010000000000000000001100
 
 
+def test_dirty_pages_invalid_address():
+    # Failed to translate
+    get_dirty_page_bitmap(addr=0xdeadbeef, expect=errno.ENOENT)
+
+    # Does not exactly match a region
+    get_dirty_page_bitmap(addr=(0x10 << PAGE_SHIFT) + 1,
+                          length=(0x20 << PAGE_SHIFT) - 1,
+                          expect=errno.ENOTSUP)
+
+    # Invalid requested bitmap size
+    get_dirty_page_bitmap(page_size=1 << 24, expect=errno.EINVAL)
+
+    # Region not mapped
+    get_dirty_page_bitmap(addr=0x40 << PAGE_SHIFT, expect=errno.EINVAL)
+
+
 def stop_logging(addr=None, length=None):
     if addr is not None:
         ranges = vfio_user_device_feature_dma_logging_range(
diff --git a/test/py/test_migration.py b/test/py/test_migration.py
index 0a59bd3..0b36e1c 100644
--- a/test/py/test_migration.py
+++ b/test/py/test_migration.py
@@ -38,6 +38,7 @@ current_state = None
 path = []
 read_data = None
 write_data = None
+fail_callbacks = False
 
 
 UNREACHABLE_STATES = {
@@ -51,6 +52,9 @@ UNREACHABLE_STATES = {
 def migr_trans_cb(_ctx, state):
     global current_state, path
 
+    if fail_callbacks:
+        return -1
+
     if state == VFU_MIGR_STATE_STOP:
         state = VFIO_USER_DEVICE_STATE_STOP
     elif state == VFU_MIGR_STATE_RUNNING:
@@ -74,17 +78,27 @@ def migr_trans_cb(_ctx, state):
 @read_data_cb_t
 def migr_read_data_cb(_ctx, buf, count):
     global read_data
+
+    if fail_callbacks:
+        return -1
+
     length = min(count, len(read_data))
     ctypes.memmove(buf, read_data, length)
     read_data = None
+
     return length
 
 
 @write_data_cb_t
 def migr_write_data_cb(_ctx, buf, count):
     global write_data
+
+    if fail_callbacks:
+        return -1
+
     write_data = bytes(count)
     ctypes.memmove(write_data, buf, count)
+
     return count
 
 
@@ -106,12 +120,16 @@ def test_migration_setup():
     assert ctx is not None
 
     cbs = vfu_migration_callbacks_t()
-    cbs.version = VFU_MIGR_CALLBACKS_VERS
+    cbs.version = 1
     cbs.transition = migr_trans_cb
     cbs.read_data = migr_read_data_cb
     cbs.write_data = migr_write_data_cb
 
     ret = vfu_setup_device_migration_callbacks(ctx, cbs)
+    assert ret < 0
+
+    cbs.version = VFU_MIGR_CALLBACKS_VERS
+    ret = vfu_setup_device_migration_callbacks(ctx, cbs)
     assert ret == 0
 
     vfu_setup_device_quiesce_cb(ctx)
@@ -232,6 +250,28 @@ def test_migration_nonexistent_state():
     transition_to_state(0xabcd, expect=errno.EINVAL)
 
 
+def test_migration_failed_callback():
+    global fail_callbacks
+    fail_callbacks = True
+    transition_to_state(VFIO_USER_DEVICE_STATE_RUNNING, expect=errno.EINVAL)
+    fail_callbacks = False
+
+
+def test_migration_get_state():
+    transition_to_state(VFIO_USER_DEVICE_STATE_RUNNING)
+
+    feature = vfio_user_device_feature(
+        argsz=len(vfio_user_device_feature()) +
+            len(vfio_user_device_feature_mig_state()),
+        flags=VFIO_DEVICE_FEATURE_GET | VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE
+    )
+
+    result = msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, feature)
+    _, result = vfio_user_device_feature.pop_from_buffer(result)
+    state, _ = vfio_user_device_feature_mig_state.pop_from_buffer(result)
+    assert state.device_state == VFIO_USER_DEVICE_STATE_RUNNING
+
+
 def test_handle_mig_data_read():
     global read_data
 
@@ -297,6 +337,29 @@ def test_handle_mig_data_read_invalid_state():
     msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload, expect=errno.EINVAL)
 
 
+def test_handle_mig_data_read_failed_callback():
+    global fail_callbacks
+
+    transition_to_state(VFIO_USER_DEVICE_STATE_PRE_COPY)
+
+    fail_callbacks = True
+
+    payload = vfio_user_mig_data(
+        argsz=len(vfio_user_mig_data()),
+        size=4
+    )
+
+    msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload, expect=errno.EINVAL)
+
+    fail_callbacks = False
+
+
+def test_handle_mig_data_read_short_write():
+    payload = struct.pack("I", 8)
+
+    msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload, expect=errno.EINVAL)
+
+
 def test_handle_mig_data_write():
     payload = vfio_user_mig_data(
         argsz=len(vfio_user_mig_data()) + 4,
@@ -310,6 +373,22 @@ def test_handle_mig_data_write():
     assert write_data == data
 
 
+def test_handle_mig_data_write_too_long():
+    payload = vfio_user_mig_data(
+        argsz=len(vfio_user_mig_data()) + 8,
+        size=8
+    )
+
+    # When we set up the tests at the top of this file we specify that the max
+    # data transfer size is 4 bytes. Here we test to check that a transfer of 8
+    # bytes fails.
+
+    data = bytes([1, 2, 3, 4, 5, 6, 7, 8])
+    transition_to_state(VFIO_USER_DEVICE_STATE_RESUMING)
+    msg(ctx, sock, VFIO_USER_MIG_DATA_WRITE, bytes(payload) + data,
+        expect=errno.EINVAL)
+
+
 def test_handle_mig_data_write_invalid_state():
     payload = vfio_user_mig_data(
         argsz=len(vfio_user_mig_data()) + 4,
@@ -323,12 +402,78 @@ def test_handle_mig_data_write_invalid_state():
         expect=errno.EINVAL)
 
 
+def test_handle_mig_data_write_failed_callback():
+    global fail_callbacks
+
+    transition_to_state(VFIO_USER_DEVICE_STATE_RESUMING)
+
+    fail_callbacks = True
+
+    payload = vfio_user_mig_data(
+        argsz=len(vfio_user_mig_data()) + 4,
+        size=4
+    )
+
+    data = bytes([1, 2, 3, 4])
+
+    msg(ctx, sock, VFIO_USER_MIG_DATA_WRITE, bytes(payload) + data,
+        expect=errno.EINVAL)
+
+    fail_callbacks = False
+
+
+def test_handle_mig_data_write_short_write():
+    payload = struct.pack("I", 8)
+
+    msg(ctx, sock, VFIO_USER_MIG_DATA_WRITE, payload, expect=errno.EINVAL)
+
+
+def test_device_feature_migration():
+    payload = vfio_user_device_feature(
+        argsz=len(vfio_user_device_feature()),
+        flags=VFIO_DEVICE_FEATURE_GET | VFIO_DEVICE_FEATURE_MIGRATION
+    )
+
+    result = msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, payload)
+    _, result = vfio_user_device_feature.pop_from_buffer(result)
+    flags, = struct.unpack("Q", result)
+
+    assert flags == VFIO_MIGRATION_STOP_COPY | VFIO_MIGRATION_PRE_COPY
+
+
 def test_device_feature_short_write():
     payload = struct.pack("I", 8)
 
     msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, payload, expect=errno.EINVAL)
 
 
+def test_device_feature_unsupported_operation():
+    payload = vfio_user_device_feature(
+        argsz=len(vfio_user_device_feature()),
+        flags=VFIO_DEVICE_FEATURE_SET | VFIO_DEVICE_FEATURE_MIGRATION
+    )
+
+    msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, payload, expect=errno.EINVAL)
+
+
+def test_device_feature_probe():
+    payload = vfio_user_device_feature(
+        argsz=len(vfio_user_device_feature()),
+        flags=VFIO_DEVICE_FEATURE_PROBE | VFIO_DEVICE_FEATURE_MIGRATION
+    )
+
+    result = msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, payload)
+    assert bytes(payload) == result
+
+    payload = vfio_user_device_feature(
+        argsz=len(vfio_user_device_feature()),
+        flags=VFIO_DEVICE_FEATURE_PROBE | VFIO_DEVICE_FEATURE_SET |
+              VFIO_DEVICE_FEATURE_GET | VFIO_DEVICE_FEATURE_MIGRATION
+    )
+
+    msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, payload, expect=errno.EINVAL)
+
+
 def test_migration_cleanup():
     disconnect_client(ctx, sock)
     vfu_destroy_ctx(ctx)