respond to more of Thanos's comments

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

6 files changed, 310 insertions, 207 deletions

diff --git a/test/py/libvfio_user.py b/test/py/libvfio_user.py
index 233f42a..e93bc4f 100644
--- a/test/py/libvfio_user.py
+++ b/test/py/libvfio_user.py
@@ -570,6 +570,13 @@ class vfio_user_device_feature(Structure):
     ]
 
 
+class vfio_user_device_feature_migration(Structure):
+    _pack_ = 1
+    _fields_ = [
+        ("flags", c.c_uint64)
+    ]
+
+
 class vfio_user_device_feature_mig_state(Structure):
     _pack_ = 1
     _fields_ = [
@@ -727,15 +734,20 @@ def connect_sock():
 def connect_client(ctx, max_data_xfer_size=None):
     sock = connect_sock()
 
-    if max_data_xfer_size is None:
-        json = b'{ "capabilities": { "max_msg_fds": 8 } }'
-    else:
-        json = b'{ "capabilities": { "max_msg_fds": 8, "max_data_xfer_size": '\
-            + str(max_data_xfer_size).encode("utf-8") + b' } }'
+    caps = {
+        "capabilities": {
+            "max_msg_fds": 8
+        }
+    }
+
+    if max_data_xfer_size is not None:
+        caps["capabilities"]["max_data_xfer_size"] = max_data_xfer_size
+
+    caps_json = json.dumps(caps).encode("utf-8")
 
     # struct vfio_user_version
-    payload = struct.pack("HH%dsc" % len(json), LIBVFIO_USER_MAJOR,
-                          LIBVFIO_USER_MINOR, json, b'\0')
+    payload = struct.pack("HH%dsc" % len(caps_json), LIBVFIO_USER_MAJOR,
+                          LIBVFIO_USER_MINOR, caps_json, b'\0')
     hdr = vfio_user_header(VFIO_USER_VERSION, size=len(payload))
     sock.send(hdr + payload)
     vfu_attach_ctx(ctx, expect=0)
@@ -989,6 +1001,18 @@ def prepare_ctx_for_dma(dma_register=__dma_register,
 
     return ctx
 
+
+def transition_to_state(ctx, sock, state, expect=0, rsp=True, busy=False):
+    feature = vfio_user_device_feature(
+        argsz=len(vfio_user_device_feature()) +
+            len(vfio_user_device_feature_mig_state()),
+        flags=VFIO_DEVICE_FEATURE_SET | VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE
+    )
+    payload = vfio_user_device_feature_mig_state(device_state=state)
+    msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, bytes(feature) + bytes(payload),
+        expect=expect, rsp=rsp, busy=busy)
+
+
 #
 # Library wrappers
 #
@@ -1254,6 +1278,11 @@ def fds_are_same(fd1: int, fd2: int) -> bool:
 
 
 def get_bitmap_size(size: int, pgsize: int) -> int:
+    """
+    Returns the size, in bytes, of the bitmap that represents the given range
+    with the given page size.
+    """
+
     nr_pages = (size // pgsize) + (1 if size % pgsize != 0 else 0)
     return ((nr_pages + 63) & ~63) // 8
 
diff --git a/test/py/test_dirty_pages.py b/test/py/test_dirty_pages.py
index 30fd757..8dfb45b 100644
--- a/test/py/test_dirty_pages.py
+++ b/test/py/test_dirty_pages.py
@@ -69,9 +69,6 @@ def test_dirty_pages_setup():
     ret = vfu_setup_device_dma(ctx, dma_register, dma_unregister)
     assert ret == 0
 
-    f = tempfile.TemporaryFile()
-    f.truncate(2 << PAGE_SHIFT)
-
     ret = vfu_realize_ctx(ctx)
     assert ret == 0
 
@@ -92,15 +89,6 @@ def test_dirty_pages_setup():
 
     msg(ctx, sock, VFIO_USER_DMA_MAP, payload)
 
-    f2 = tempfile.TemporaryFile()
-    f2.truncate(0x10 << PAGE_SHIFT)
-
-    payload = vfio_user_dma_map(argsz=len(vfio_user_dma_map()),
-        flags=(VFIO_USER_F_DMA_REGION_READ | VFIO_USER_F_DMA_REGION_WRITE),
-        offset=0, addr=0x60 << PAGE_SHIFT, size=0x20 << PAGE_SHIFT)
-
-    msg(ctx, sock, VFIO_USER_DMA_MAP, payload, fds=[f2.fileno()])
-
 
 def test_setup_migration():
     ret = vfu_setup_device_migration_callbacks(ctx)
@@ -108,13 +96,27 @@ def test_setup_migration():
 
 
 def start_logging(addr=None, length=None, page_size=PAGE_SIZE, expect=0):
+    """
+    Start logging dirty writes.
+
+    If a region and page size are specified, they will be sent to the server to
+    start logging. Otherwise, all regions will be logged and the default page
+    size will be used.
+
+    Note: in the current implementation, all regions are logged whether or not
+    you specify a region, as the additional constraint of only logging a
+    certain region is considered an optimisation and is not yet implemented.
+    """
+
     if addr is not None:
         ranges = vfio_user_device_feature_dma_logging_range(
             iova=addr,
             length=length
         )
+        num_ranges = 1
     else:
-        ranges = []
+        ranges = bytearray()
+        num_ranges = 0
 
     feature = vfio_user_device_feature(
         argsz=len(vfio_user_device_feature()) +
@@ -124,13 +126,17 @@ def start_logging(addr=None, length=None, page_size=PAGE_SIZE, expect=0):
 
     payload = vfio_user_device_feature_dma_logging_control(
         page_size=page_size,
-        num_ranges=(1 if addr is not None else 0),
+        num_ranges=num_ranges,
         reserved=0)
 
     msg(ctx, sock, VFIO_USER_DEVICE_FEATURE,
         bytes(feature) + bytes(payload) + bytes(ranges), expect=expect)
 
 
+def test_dirty_pages_start_zero_pgsize():
+    start_logging(page_size=0, expect=errno.EINVAL)
+
+
 def test_dirty_pages_start():
     start_logging()
     # should be idempotent
@@ -138,55 +144,27 @@ def test_dirty_pages_start():
 
 
 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()) + \
-            8  # size of bitmap
-
-    feature = vfio_user_device_feature(
-        argsz=argsz,
-        flags=VFIO_DEVICE_FEATURE_DMA_LOGGING_REPORT | VFIO_DEVICE_FEATURE_GET
-    )
-
-    report = vfio_user_device_feature_dma_logging_report(
-        iova=0x10 << PAGE_SHIFT,
-        length=0x10 << PAGE_SHIFT,
-        page_size=PAGE_SIZE
-    )
-
-    payload = bytes(feature) + bytes(report)
-
-    result = msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, payload)
-
-    assert len(result) == argsz
+    """
+    Once we've started logging with page size PAGE_SIZE, any request to start
+    logging at a different page size should be rejected.
+    """
 
-    feature, result = vfio_user_device_feature.pop_from_buffer(result)
-
-    assert feature.argsz == argsz
-    assert feature.flags == VFIO_DEVICE_FEATURE_DMA_LOGGING_REPORT \
-        | VFIO_DEVICE_FEATURE_GET
-
-    report, bitmap = \
-        vfio_user_device_feature_dma_logging_report.pop_from_buffer(result)
-
-    assert report.iova == 0x10 << PAGE_SHIFT
-    assert report.length == 0x10 << PAGE_SHIFT
-    assert report.page_size == PAGE_SIZE
-
-    assert len(bitmap) == 8
-
-    for b in bitmap:
-        assert b == 0
+    start_logging(page_size=PAGE_SIZE >> 1, expect=errno.EINVAL)
+    start_logging(page_size=PAGE_SIZE << 1, expect=errno.EINVAL)
 
 
 def get_dirty_page_bitmap(addr=0x10 << PAGE_SHIFT, length=0x10 << PAGE_SHIFT,
                           page_size=PAGE_SIZE, expect=0):
+    """
+    Get the dirty page bitmap from the server for the given region and page
+    size as a 64-bit integer. This only works for bitmaps that fit within a
+    64-bit integer.
+    """
+
     bitmap_size = get_bitmap_size(length, page_size)
 
+    assert bitmap_size == 8
+
     argsz = len(vfio_user_device_feature()) + \
             len(vfio_user_device_feature_dma_logging_report()) + \
             bitmap_size
@@ -220,6 +198,11 @@ def get_dirty_page_bitmap(addr=0x10 << PAGE_SHIFT, length=0x10 << PAGE_SHIFT,
     return struct.unpack("Q", result)[0]
 
 
+def test_dirty_pages_get_unmodified():
+    bitmap = get_dirty_page_bitmap()
+    assert bitmap == 0
+
+
 sg3 = None
 iovec3 = None
 
@@ -354,11 +337,11 @@ def test_dirty_pages_get_modified():
     assert bitmap == 0b010000000000000000001100
 
 
-def test_dirty_pages_invalid_address():
+def test_dirty_pages_invalid_arguments():
     # Failed to translate
     get_dirty_page_bitmap(addr=0xdeadbeef, expect=errno.ENOENT)
 
-    # Does not exactly match a region
+    # Does not exactly match a region (libvfio-user limitation)
     get_dirty_page_bitmap(addr=(0x10 << PAGE_SHIFT) + 1,
                           length=(0x20 << PAGE_SHIFT) - 1,
                           expect=errno.ENOTSUP)
@@ -402,4 +385,28 @@ def test_dirty_pages_cleanup():
     disconnect_client(ctx, sock)
     vfu_destroy_ctx(ctx)
 
+
+def test_dirty_pages_uninitialised_dma():
+    global ctx, sock
+
+    ctx = vfu_create_ctx(flags=LIBVFIO_USER_FLAG_ATTACH_NB)
+    assert ctx is not None
+
+    ret = vfu_pci_init(ctx)
+    assert ret == 0
+
+    vfu_setup_device_quiesce_cb(ctx, quiesce_cb=quiesce_cb)
+
+    ret = vfu_realize_ctx(ctx)
+    assert ret == 0
+
+    sock = connect_client(ctx)
+
+    start_logging(expect=errno.EINVAL)
+    get_dirty_page_bitmap(expect=errno.EINVAL)
+
+    disconnect_client(ctx, sock)
+
+    vfu_destroy_ctx(ctx)
+
 # ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
diff --git a/test/py/test_migration.py b/test/py/test_migration.py
index f8a1eee..674945c 100644
--- a/test/py/test_migration.py
+++ b/test/py/test_migration.py
@@ -40,6 +40,16 @@ path = []
 read_data = None
 write_data = None
 fail_callbacks = False
+fail_callbacks_errno = None
+
+
+STATES = {
+    VFIO_USER_DEVICE_STATE_STOP,
+    VFIO_USER_DEVICE_STATE_RUNNING,
+    VFIO_USER_DEVICE_STATE_STOP_COPY,
+    VFIO_USER_DEVICE_STATE_RESUMING,
+    VFIO_USER_DEVICE_STATE_PRE_COPY
+}
 
 
 UNREACHABLE_STATES = {
@@ -49,23 +59,25 @@ UNREACHABLE_STATES = {
 }
 
 
+VFU_TO_VFIO_MIGR_STATE = {
+    VFU_MIGR_STATE_STOP: VFIO_USER_DEVICE_STATE_STOP,
+    VFU_MIGR_STATE_RUNNING: VFIO_USER_DEVICE_STATE_RUNNING,
+    VFU_MIGR_STATE_STOP_AND_COPY: VFIO_USER_DEVICE_STATE_STOP_COPY,
+    VFU_MIGR_STATE_RESUME: VFIO_USER_DEVICE_STATE_RESUMING,
+    VFU_MIGR_STATE_PRE_COPY: VFIO_USER_DEVICE_STATE_PRE_COPY
+}
+
+
 @transition_cb_t
 def migr_trans_cb(_ctx, state):
     global current_state, path
 
     if fail_callbacks:
+        c.set_errno(fail_callbacks_errno)
         return -1
 
-    if state == VFU_MIGR_STATE_STOP:
-        state = VFIO_USER_DEVICE_STATE_STOP
-    elif state == VFU_MIGR_STATE_RUNNING:
-        state = VFIO_USER_DEVICE_STATE_RUNNING
-    elif state == VFU_MIGR_STATE_STOP_AND_COPY:
-        state = VFIO_USER_DEVICE_STATE_STOP_COPY
-    elif state == VFU_MIGR_STATE_RESUME:
-        state = VFIO_USER_DEVICE_STATE_RESUMING
-    elif state == VFU_MIGR_STATE_PRE_COPY:
-        state = VFIO_USER_DEVICE_STATE_PRE_COPY
+    if state in VFU_TO_VFIO_MIGR_STATE:
+        state = VFU_TO_VFIO_MIGR_STATE[state]
     else:
         assert False
 
@@ -103,15 +115,17 @@ def migr_write_data_cb(_ctx, buf, count):
     return count
 
 
-def transition_to_state(state, expect=0):
-    feature = vfio_user_device_feature(
-        argsz=len(vfio_user_device_feature()) +
-            len(vfio_user_device_feature_mig_state()),
-        flags=VFIO_DEVICE_FEATURE_SET | VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE
-    )
-    payload = vfio_user_device_feature_mig_state(device_state=state)
-    msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, bytes(feature) + bytes(payload),
-        expect=expect)
+def setup_fail_callbacks(errno=errno.EINVAL):
+    global fail_callbacks, fail_callbacks_errno
+    fail_callbacks = True
+    fail_callbacks_errno = errno
+
+
+def teardown_fail_callbacks():
+    global fail_callbacks, fail_callbacks_errno
+    fail_callbacks = False
+    fail_callbacks_errno = None
+    c.set_errno(0)
 
 
 def test_migration_setup():
@@ -121,15 +135,15 @@ def test_migration_setup():
     assert ctx is not None
 
     cbs = vfu_migration_callbacks_t()
-    cbs.version = 1
+    cbs.version = 1  # old callbacks version
     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
+    assert ret < 0, "do not allow old callbacks version"
 
-    cbs.version = VFU_MIGR_CALLBACKS_VERS
+    cbs.version = VFU_MIGR_CALLBACKS_VERS  # new callbacks version
     ret = vfu_setup_device_migration_callbacks(ctx, cbs)
     assert ret == 0
 
@@ -138,26 +152,32 @@ def test_migration_setup():
     ret = vfu_realize_ctx(ctx)
     assert ret == 0
 
-    sock = connect_client(ctx, 4)
+    max_data_xfer_size = 4  # for later tests
+    sock = connect_client(ctx, max_data_xfer_size)
 
 
 def get_server_shortest_path(a, b, expectA=0, expectB=0):
+    """
+    Carry out the state transition from a to b on the server, keeping track of
+    and returning the transition path taken.
+    """
+
     global path
 
     if current_state == VFIO_USER_DEVICE_STATE_STOP_COPY and \
             a == VFIO_USER_DEVICE_STATE_PRE_COPY:
         # The transition STOP_COPY -> PRE_COPY is explicitly blocked so we
         # advance one state to get around this in order to set up the test.
-        transition_to_state(VFIO_USER_DEVICE_STATE_STOP)
+        transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_STOP)
 
-    transition_to_state(a, expect=expectA)
+    transition_to_state(ctx, sock, a, expect=expectA)
 
     if expectA != 0:
         return None
 
     path = []
 
-    transition_to_state(b, expect=expectB)
+    transition_to_state(ctx, sock, b, expect=expectB)
 
     return path.copy()
 
@@ -176,12 +196,6 @@ def test_migration_shortest_state_transition_paths():
     by the implementation correctly follow these rules.
     """
 
-    # states (vertices)
-    V = {VFIO_USER_DEVICE_STATE_ERROR, VFIO_USER_DEVICE_STATE_STOP,
-         VFIO_USER_DEVICE_STATE_RUNNING, VFIO_USER_DEVICE_STATE_STOP_COPY,
-         VFIO_USER_DEVICE_STATE_RESUMING, VFIO_USER_DEVICE_STATE_RUNNING_P2P,
-         VFIO_USER_DEVICE_STATE_PRE_COPY, VFIO_USER_DEVICE_STATE_PRE_COPY_P2P}
-
     # allowed direct transitions (edges)
     E = {
         VFIO_USER_DEVICE_STATE_ERROR: set(),
@@ -204,26 +218,40 @@ def test_migration_shortest_state_transition_paths():
         VFIO_USER_DEVICE_STATE_PRE_COPY_P2P: set()
     }
 
+    # states (vertices)
+    V = E.keys()
+
     # "saving states" which cannot be internal arcs
-    S = {VFIO_USER_DEVICE_STATE_PRE_COPY, VFIO_USER_DEVICE_STATE_STOP_COPY}
+    saving_states = {VFIO_USER_DEVICE_STATE_PRE_COPY,
+                     VFIO_USER_DEVICE_STATE_STOP_COPY}
 
+    # Consider each vertex in turn to be the start state, that is, the state
+    # we are transitioning from.
     for source in V:
         back = {v: None for v in V}
         queue = deque([(source, None)])
 
+        # Use BFS to calculate the shortest path from the start state to every
+        # other state, following the rule that no intermediate states can be
+        # saving states.
         while len(queue) > 0:
             (curr, prev) = queue.popleft()
             back[curr] = prev
             for nxt in E[curr]:
-                if back[nxt] is None and (curr == source or curr not in S):
+                if back[nxt] is None \
+                        and (curr == source or curr not in saving_states):
                     queue.append((nxt, curr))
 
+        # Iterate over the states
         for target in V:
             if source == VFIO_USER_DEVICE_STATE_STOP_COPY \
                     and target == VFIO_USER_DEVICE_STATE_PRE_COPY:
                 # test for this transition being blocked in a separate test
                 continue
 
+            # If BFS found a path to that state, follow the backpointers to
+            # calculate the path, and check that it's equal to the path taken
+            # by the server.
             if back[target] is not None:
                 seq = deque([])
                 curr = target
@@ -235,31 +263,41 @@ def test_migration_shortest_state_transition_paths():
 
                 assert len(seq) == len(server_seq)
                 assert all(seq[i] == server_seq[i] for i in range(len(seq)))
+
+            # If BFS couldn't find a path to that state, check that the server
+            # doesn't allow that transition either.
             else:
-                expectA = 22 if source in UNREACHABLE_STATES else 0
+                # If the start state is an unreachable state, we won't be able
+                # to transition into it in order to try and calculate a path on
+                # the server, so we expect that transition to fail.
+                expectA = errno.EINVAL if source in UNREACHABLE_STATES else 0
 
+                # No matter what, we expect transitioning to the target state
+                # to fail.
                 get_server_shortest_path(source, target, expectA=expectA,
-                                         expectB=22)
+                                         expectB=errno.EINVAL)
 
 
-def test_migration_stop_copy_to_pre_copy_blocked():
-    transition_to_state(VFIO_USER_DEVICE_STATE_STOP_COPY)
-    transition_to_state(VFIO_USER_DEVICE_STATE_PRE_COPY, expect=errno.EINVAL)
+def test_migration_stop_copy_to_pre_copy_rejected():
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_STOP_COPY)
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_PRE_COPY,
+                        expect=errno.EINVAL)
 
 
 def test_migration_nonexistent_state():
-    transition_to_state(0xabcd, expect=errno.EINVAL)
+    transition_to_state(ctx, sock, 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
+    setup_fail_callbacks()
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_RUNNING,
+                        expect=errno.EINVAL)
+    assert c.get_errno() == errno.EINVAL
+    teardown_fail_callbacks()
 
 
 def test_migration_get_state():
-    transition_to_state(VFIO_USER_DEVICE_STATE_RUNNING)
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_RUNNING)
 
     feature = vfio_user_device_feature(
         argsz=len(vfio_user_device_feature()) +
@@ -276,181 +314,193 @@ def test_migration_get_state():
 def test_handle_mig_data_read():
     global read_data
 
-    transition_to_state(VFIO_USER_DEVICE_STATE_RUNNING)
-
-    argsz = len(vfio_user_mig_data()) + 4
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_RUNNING)
 
+    data = bytes([0, 1, 2, 3])
+    argsz = len(vfio_user_mig_data()) + len(data)
     payload = vfio_user_mig_data(
         argsz=argsz,
-        size=4
+        size=len(data)
     )
 
-    data = bytes([0, 1, 2, 3])
+    VALID_STATES = {VFIO_USER_DEVICE_STATE_PRE_COPY,
+                    VFIO_USER_DEVICE_STATE_STOP_COPY}
 
-    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
-    assert result[len(vfio_user_mig_data()):] == data
+    for state in STATES:
+        transition_to_state(ctx, sock, state)
+        read_data = data
+        expect = 0 if state in VALID_STATES else errno.EINVAL
+        result = msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload,
+                     expect=expect)
 
-    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
-    assert result[len(vfio_user_mig_data()):] == data
+        if state in VALID_STATES:
+            assert len(result) == argsz
+            assert result[len(vfio_user_mig_data()):] == data
 
 
 def test_handle_mig_data_read_too_long():
-    global read_data
-
-    transition_to_state(VFIO_USER_DEVICE_STATE_RUNNING)
+    """
+    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.
+    """
 
-    payload = vfio_user_mig_data(
-        argsz=len(vfio_user_mig_data()) + 8,
-        size=8
-    )
+    global read_data
 
-    # 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.
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_RUNNING)
 
-    transition_to_state(VFIO_USER_DEVICE_STATE_PRE_COPY)
     read_data = bytes([1, 2, 3, 4, 5, 6, 7, 8])
-    msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload, expect=errno.EINVAL)
-
-
-def test_handle_mig_data_read_invalid_state():
-    global read_data
 
     payload = vfio_user_mig_data(
-        argsz=len(vfio_user_mig_data()) + 4,
-        size=4
+        argsz=len(vfio_user_mig_data()) + len(read_data),
+        size=len(read_data)
     )
 
-    data = bytes([1, 2, 3, 4])
-
-    transition_to_state(VFIO_USER_DEVICE_STATE_RUNNING)
-    read_data = data
-    msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload, expect=errno.EINVAL)
-
-    transition_to_state(VFIO_USER_DEVICE_STATE_STOP)
-    read_data = data
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_PRE_COPY)
     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(ctx, sock, VFIO_USER_DEVICE_STATE_PRE_COPY)
 
-    transition_to_state(VFIO_USER_DEVICE_STATE_PRE_COPY)
+    read_data = bytes([1, 2, 3, 4])
 
-    fail_callbacks = True
+    setup_fail_callbacks()
 
     payload = vfio_user_mig_data(
-        argsz=len(vfio_user_mig_data()) + 4,
-        size=4
+        argsz=len(vfio_user_mig_data()) + len(read_data),
+        size=len(read_data)
     )
 
     msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload, expect=errno.EINVAL)
+    assert c.get_errno() == errno.EINVAL
 
-    fail_callbacks = False
+    teardown_fail_callbacks()
 
 
 def test_handle_mig_data_read_short_write():
-    payload = struct.pack("I", 8)
+    data = bytes([1, 2, 3, 4])
 
-    msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload, expect=errno.EINVAL)
+    payload = vfio_user_mig_data(
+        argsz=len(vfio_user_mig_data()) + len(data),
+        size=len(data)
+    )
+
+    payload = bytes(payload)
+    assert len(payload) == 8
+
+    # don't send the last byte
+    msg(ctx, sock, VFIO_USER_MIG_DATA_READ, payload[0:7], expect=errno.EINVAL)
 
 
 def test_handle_mig_data_write():
+    data = bytes([1, 2, 3, 4])
+
     payload = vfio_user_mig_data(
-        argsz=len(vfio_user_mig_data()) + 4,
-        size=4
+        argsz=len(vfio_user_mig_data()) + len(data),
+        size=len(data)
     )
 
-    data = bytes([1, 2, 3, 4])
-
-    transition_to_state(VFIO_USER_DEVICE_STATE_RESUMING)
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_RESUMING)
     msg(ctx, sock, VFIO_USER_MIG_DATA_WRITE, bytes(payload) + data)
     assert write_data == data
 
 
-def test_handle_mig_data_write_too_long():
+def test_handle_mig_data_write_invalid_state():
+    data = bytes([1, 2, 3, 4])
+
     payload = vfio_user_mig_data(
-        argsz=len(vfio_user_mig_data()) + 8,
-        size=8
+        argsz=len(vfio_user_mig_data()) + len(data),
+        size=len(data)
     )
 
-    # 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)
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_RUNNING)
     msg(ctx, sock, VFIO_USER_MIG_DATA_WRITE, bytes(payload) + data,
         expect=errno.EINVAL)
 
 
-def test_handle_mig_data_write_invalid_state():
+def test_handle_mig_data_write_too_long():
+    """
+    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])
+
     payload = vfio_user_mig_data(
-        argsz=len(vfio_user_mig_data()) + 4,
-        size=4
+        argsz=len(vfio_user_mig_data()) + len(data),
+        size=len(data)
     )
 
-    data = bytes([1, 2, 3, 4])
-
-    transition_to_state(VFIO_USER_DEVICE_STATE_RUNNING)
+    transition_to_state(ctx, sock, 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_failed_callback():
-    global fail_callbacks
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_RESUMING)
 
-    transition_to_state(VFIO_USER_DEVICE_STATE_RESUMING)
+    data = bytes([1, 2, 3, 4])
 
-    fail_callbacks = True
+    setup_fail_callbacks()
 
     payload = vfio_user_mig_data(
-        argsz=len(vfio_user_mig_data()) + 4,
-        size=4
+        argsz=len(vfio_user_mig_data()) + len(data),
+        size=len(data)
     )
 
-    data = bytes([1, 2, 3, 4])
-
     msg(ctx, sock, VFIO_USER_MIG_DATA_WRITE, bytes(payload) + data,
         expect=errno.EINVAL)
+    assert c.get_errno() == errno.EINVAL
 
-    fail_callbacks = False
+    teardown_fail_callbacks()
 
 
 def test_handle_mig_data_write_short_write():
-    payload = struct.pack("I", 8)
+    data = bytes([1, 2, 3, 4])
+
+    payload = vfio_user_mig_data(
+        argsz=len(vfio_user_mig_data()) + len(data),
+        size=len(data)
+    )
 
     msg(ctx, sock, VFIO_USER_MIG_DATA_WRITE, payload, expect=errno.EINVAL)
 
 
-def test_device_feature_migration():
+def test_device_feature_migration_get():
     payload = vfio_user_device_feature(
-        argsz=len(vfio_user_device_feature()) + 8,
+        argsz=len(vfio_user_device_feature()) +
+              len(vfio_user_device_feature_migration()),
         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)
+    flags, _ = vfio_user_device_feature_migration.pop_from_buffer(result)
+    flags = flags.flags
 
     assert flags == VFIO_MIGRATION_STOP_COPY | VFIO_MIGRATION_PRE_COPY
 
 
 def test_device_feature_short_write():
-    payload = struct.pack("I", 8)
+    payload = vfio_user_device_feature(
+        argsz=len(vfio_user_device_feature()) +
+              len(vfio_user_device_feature_migration()),
+        flags=VFIO_DEVICE_FEATURE_GET | VFIO_DEVICE_FEATURE_MIGRATION
+    )
 
-    msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, payload, expect=errno.EINVAL)
+    payload = bytes(payload)
+    assert len(payload) == 8
+
+    # don't send the last byte
+    msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, payload[0:7], expect=errno.EINVAL)
 
 
 def test_device_feature_unsupported_operation():
     payload = vfio_user_device_feature(
-        argsz=len(vfio_user_device_feature()) + 8,
+        argsz=len(vfio_user_device_feature()) +
+              len(vfio_user_device_feature_migration()),
         flags=VFIO_DEVICE_FEATURE_SET | VFIO_DEVICE_FEATURE_MIGRATION
     )
 
@@ -478,7 +528,7 @@ def test_device_feature_bad_argsz_get_migration():
 def test_device_feature_bad_argsz_get_dma():
     argsz = len(vfio_user_device_feature()) + \
             len(vfio_user_device_feature_dma_logging_report()) + \
-            8  # bitmap size
+            get_bitmap_size(0x20 << PAGE_SHIFT, PAGE_SIZE)
 
     feature = vfio_user_device_feature(
         argsz=argsz - 1,  # not big enough
@@ -529,4 +579,4 @@ def test_migration_cleanup():
     disconnect_client(ctx, sock)
     vfu_destroy_ctx(ctx)
 
-# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_quiesce.py b/test/py/test_quiesce.py
index e97c565..baa54cd 100644
--- a/test/py/test_quiesce.py
+++ b/test/py/test_quiesce.py
@@ -195,6 +195,31 @@ def test_allowed_funcs_in_quiesced_dma_unregister_busy(mock_quiesce,
     mock_dma_unregister.assert_called_once_with(ctx, mock.ANY)
 
 
+@patch('libvfio_user.migr_trans_cb', side_effect=_side_effect)
+@patch('libvfio_user.quiesce_cb')
+def test_allowed_funcs_in_quiesed_migration(mock_quiesce,
+                                            mock_trans):
+
+    global ctx, sock
+    _map_dma_region(ctx, sock)
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_STOP)
+    mock_trans.assert_called_once_with(ctx, VFU_MIGR_STATE_STOP)
+
+
+@patch('libvfio_user.migr_trans_cb', side_effect=_side_effect)
+@patch('libvfio_user.quiesce_cb')
+def test_allowed_funcs_in_quiesed_migration_busy(mock_quiesce,
+                                                 mock_trans):
+
+    global ctx, sock
+    _map_dma_region(ctx, sock)
+    mock_quiesce.side_effect = fail_with_errno(errno.EBUSY)
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_STOP, busy=True)
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
+    mock_trans.assert_called_once_with(ctx, VFU_MIGR_STATE_STOP)
+
+
 @patch('libvfio_user.reset_cb', side_effect=_side_effect)
 @patch('libvfio_user.quiesce_cb')
 def test_allowed_funcs_in_quiesced_reset(mock_quiesce, mock_reset):
diff --git a/test/py/test_request_errors.py b/test/py/test_request_errors.py
index a82174c..ee7b4df 100644
--- a/test/py/test_request_errors.py
+++ b/test/py/test_request_errors.py
@@ -198,16 +198,8 @@ def test_reply_fail_quiesce_busy(mock_migr_trans_cb, mock_quiesce,
     mock_migr_trans_cb.side_effect = migr_trans_cb_side_effect
 
     # change the state, it should close the socket causing the reply to fail
-    feature = vfio_user_device_feature(
-        argsz=len(vfio_user_device_feature()) +
-            len(vfio_user_device_feature_mig_state()),
-        flags=VFIO_DEVICE_FEATURE_SET | VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE
-    )
-    payload = vfio_user_device_feature_mig_state(
-        device_state=VFIO_USER_DEVICE_STATE_STOP_COPY
-    )
-    msg(ctx, sock, VFIO_USER_DEVICE_FEATURE, bytes(feature) + bytes(payload),
-        rsp=False, busy=True)
+    transition_to_state(ctx, sock, VFIO_USER_DEVICE_STATE_STOP_COPY, rsp=False,
+                        busy=True)
 
     # vfu_run_ctx will try to reset the context and to do that it needs to
     # quiesce the device first
diff --git a/test/unit-tests.c b/test/unit-tests.c
index 36faaf9..d0c08f0 100644
--- a/test/unit-tests.c
+++ b/test/unit-tests.c
@@ -407,7 +407,7 @@ test_device_is_stopped_and_copying(UNUSED void **state)
     size_t i;
     struct migration migration;
     vfu_ctx.migration = &migration;
-    for (i = 0; i < 8; i++) {
+    for (i = 0; i < VFIO_USER_DEVICE_NUM_STATES; i++) {
         migration.state = i;
         bool r = device_is_stopped_and_copying(vfu_ctx.migration);
         if (i == VFIO_USER_DEVICE_STATE_STOP_COPY) {