introduce device quiesce callback (#609)

Signed-off-by: Thanos Makatos <thanos.makatos@nutanix.com>
Reviewed-by: John Leon <john.levon@nutanix.com>

22 files changed, 1048 insertions, 281 deletions

diff --git a/test/mocks.c b/test/mocks.c
index 29fe298..619ccf9 100644
--- a/test/mocks.c
+++ b/test/mocks.c
@@ -300,12 +300,11 @@ mock_dma_register(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info)
     check_expected(info);
 }
 
-int
+void
 mock_dma_unregister(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info)
 {
     check_expected(vfu_ctx);
     check_expected(info);
-    return mock();
 }
 
 int
diff --git a/test/mocks.h b/test/mocks.h
index 7547956..c260e8f 100644
--- a/test/mocks.h
+++ b/test/mocks.h
@@ -36,7 +36,7 @@ void patch(const char *name);
 
 void mock_dma_register(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info);
 
-int mock_dma_unregister(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info);
+void mock_dma_unregister(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info);
 
 int mock_reset_cb(vfu_ctx_t *vfu_ctx, vfu_reset_type_t type);
 
diff --git a/test/py/libvfio_user.py b/test/py/libvfio_user.py
index cbe6156..60c5fbf 100644
--- a/test/py/libvfio_user.py
+++ b/test/py/libvfio_user.py
@@ -41,6 +41,8 @@ import os
 import socket
 import struct
 import syslog
+import copy
+import tempfile
 
 UINT64_MAX = 18446744073709551615
 
@@ -280,6 +282,23 @@ class iovec_t(Structure):
         ("iov_len", c.c_int32)
     ]
 
+    def __eq__(self, other):
+        if type(self) != type(other):
+            return False
+        return self.iov_base == other.iov_base \
+            and self.iov_len == other.iov_len
+
+    def __str__(self):
+        return "%s-%s" % \
+            (hex(self.iov_base or 0), hex((self.iov_base or 0) + self.iov_len))
+
+    def __copy__(self):
+        cls = self.__class__
+        result = cls.__new__(cls)
+        result.iov_base = self.iov_base
+        result.iov_len = self.iov_len
+        return result
+
 
 class vfio_irq_info(Structure):
     _pack_ = 1
@@ -438,6 +457,30 @@ class vfu_dma_info_t(Structure):
         ("prot", c.c_uint32)
     ]
 
+    def __eq__(self, other):
+        if type(self) != type(other):
+            return False
+        return self.iova == other.iova \
+            and self.vaddr == other.vaddr \
+            and self.mapping == other.mapping \
+            and self.page_size == other.page_size \
+            and self.prot == other.prot
+
+    def __str__(self):
+        return "IOVA=%s vaddr=%s mapping=%s page_size=%s prot=%s" % \
+            (self.iova, self.vaddr, self.mapping, hex(self.page_size),
+            bin(self.prot))
+
+    def __copy__(self):
+        cls = self.__class__
+        result = cls.__new__(cls)
+        result.iova = self.iova
+        result.vaddr = self.vaddr
+        result.mapping = self.mapping
+        result.page_size = self.page_size
+        result.prot = self.prot
+        return result
+
 
 class vfio_user_dirty_pages(Structure):
     _pack_ = 1
@@ -497,6 +540,23 @@ class dma_sg_t(Structure):
         ("le_prev", c.c_void_p),
     ]
 
+    def __str__(self):
+        return "DMA addr=%s, region index=%s, length=%s, offset=%s, RW=%s" % \
+            (hex(self.dma_addr), self.region, hex(self.length),
+                hex(self.offset), self.writeable)
+
+
+class vfio_user_migration_info(Structure):
+    _pack_ = 1
+    _fields_ = [
+        ("device_state", c.c_uint32),
+        ("reserved", c.c_uint32),
+        ("pending_bytes", c.c_uint64),
+        ("data_offset", c.c_uint64),
+        ("data_size", c.c_uint64),
+    ]
+
+
 #
 # Util functions
 #
@@ -529,10 +589,14 @@ lib.vfu_pci_find_next_capability.argtypes = (c.c_void_p, c.c_bool, c.c_ulong,
                                              c.c_int)
 lib.vfu_pci_find_next_capability.restype = (c.c_ulong)
 lib.vfu_irq_trigger.argtypes = (c.c_void_p, c.c_uint)
+vfu_device_quiesce_cb_t = c.CFUNCTYPE(c.c_int, c.c_void_p, use_errno=True)
+lib.vfu_setup_device_quiesce_cb.argtypes = (c.c_void_p,
+                                            vfu_device_quiesce_cb_t)
 vfu_dma_register_cb_t = c.CFUNCTYPE(None, c.c_void_p,
-                                    c.POINTER(vfu_dma_info_t))
-vfu_dma_unregister_cb_t = c.CFUNCTYPE(c.c_int, c.c_void_p,
-                                      c.POINTER(vfu_dma_info_t))
+                                    c.POINTER(vfu_dma_info_t), use_errno=True)
+vfu_dma_unregister_cb_t = c.CFUNCTYPE(None, c.c_void_p,
+                                      c.POINTER(vfu_dma_info_t),
+                                      use_errno=True)
 lib.vfu_setup_device_dma.argtypes = (c.c_void_p, vfu_dma_register_cb_t,
                                      vfu_dma_unregister_cb_t)
 lib.vfu_setup_device_migration_callbacks.argtypes = (c.c_void_p,
@@ -548,7 +612,7 @@ lib.vfu_create_ioeventfd.argtypes = (c.c_void_p, c.c_uint32, c.c_int,
                                      c.c_size_t, c.c_uint32, c.c_uint32,
                                      c.c_uint64)
 
-lib.vfu_migr_done.argtypes = (c.c_void_p, c.c_int)
+lib.vfu_device_quiesced.argtypes = (c.c_void_p, c.c_int)
 
 
 def to_byte(val):
@@ -597,7 +661,7 @@ def disconnect_client(ctx, sock):
     # notice client closed connection
     ret = vfu_run_ctx(ctx)
     assert ret == -1
-    assert c.get_errno() == errno.ENOTCONN
+    assert c.get_errno() == errno.ENOTCONN, os.strerror(c.get_errno())
 
 
 def get_reply(sock, expect=0):
@@ -608,7 +672,8 @@ def get_reply(sock, expect=0):
     return buf[16:]
 
 
-def msg(ctx, sock, cmd, payload, expect=0, fds=None, rsp=True):
+def msg(ctx, sock, cmd, payload, expect_reply_errno=0, fds=None, rsp=True,
+        expect_run_ctx_errno=None):
     """Round trip a request and reply to the server."""
     hdr = vfio_user_header(cmd, size=len(payload))
 
@@ -618,12 +683,13 @@ def msg(ctx, sock, cmd, payload, expect=0, fds=None, rsp=True):
     else:
         sock.send(hdr + payload)
 
-    ret = vfu_run_ctx(ctx)
-    assert ret >= 0
+    ret = vfu_run_ctx(ctx, expect_errno=expect_run_ctx_errno)
+    if expect_run_ctx_errno is None:
+        assert ret >= 0, os.strerror(c.get_errno())
 
     if not rsp:
         return
-    return get_reply(sock, expect=expect)
+    return get_reply(sock, expect=expect_reply_errno)
 
 
 def get_reply_fds(sock, expect=0):
@@ -698,7 +764,7 @@ def write_pci_cfg_space(ctx, buf, count, offset, extended=False):
 
 
 def access_region(ctx, sock, is_write, region, offset, count,
-                  data=None, expect=0, rsp=True):
+                  data=None, expect=0, rsp=True, expect_run_ctx_errno=None):
     # struct vfio_user_region_access
     payload = struct.pack("QII", offset, region, count)
     if is_write:
@@ -706,22 +772,26 @@ def access_region(ctx, sock, is_write, region, offset, count,
 
     cmd = VFIO_USER_REGION_WRITE if is_write else VFIO_USER_REGION_READ
 
-    result = msg(ctx, sock, cmd, payload, expect=expect, rsp=rsp)
+    result = msg(ctx, sock, cmd, payload, expect_reply_errno=expect, rsp=rsp,
+                 expect_run_ctx_errno=expect_run_ctx_errno)
 
     if is_write:
         return None
 
-    return skip("QII", result)
+    if rsp:
+        return skip("QII", result)
 
 
-def write_region(ctx, sock, region, offset, count, data, expect=0, rsp=True):
+def write_region(ctx, sock, region, offset, count, data, expect=0, rsp=True,
+                 expect_run_ctx_errno=None):
     access_region(ctx, sock, True, region, offset, count, data, expect=expect,
-                  rsp=rsp)
+                  rsp=rsp, expect_run_ctx_errno=expect_run_ctx_errno)
 
 
-def read_region(ctx, sock, region, offset, count, expect=0):
+def read_region(ctx, sock, region, offset, count, expect=0, rsp=True,
+                expect_run_ctx_errno=None):
     return access_region(ctx, sock, False, region, offset, count,
-                         expect=expect)
+        expect=expect, rsp=rsp, expect_run_ctx_errno=expect_run_ctx_errno)
 
 
 def ext_cap_hdr(buf, offset):
@@ -733,18 +803,60 @@ def ext_cap_hdr(buf, offset):
     return cap_id, cap_next
 
 
-@vfu_dma_register_cb_t
 def dma_register(ctx, info):
     pass
 
 
-@vfu_dma_unregister_cb_t
+@vfu_dma_register_cb_t
+def __dma_register(ctx, info):
+    # The copy is required because in case of deliberate failure (e.g.
+    # test_dma_map_busy_reply_fail) the memory gets deallocated and mock only
+    # records the pointer, so the contents are all null/zero.
+    dma_register(ctx, copy.copy(info.contents))
+
+
 def dma_unregister(ctx, info):
     pass
+
+
+@vfu_dma_unregister_cb_t
+def __dma_unregister(ctx, info):
+    dma_unregister(ctx, info)
+
+
+def quiesce_cb(ctx):
     return 0
 
 
-def prepare_ctx_for_dma():
+@vfu_device_quiesce_cb_t
+def _quiesce_cb(ctx):
+    return quiesce_cb(ctx)
+
+
+def vfu_setup_device_quiesce_cb(ctx, quiesce_cb=_quiesce_cb):
+    assert ctx is not None
+    lib.vfu_setup_device_quiesce_cb(ctx,
+                                       c.cast(quiesce_cb,
+                                              vfu_device_quiesce_cb_t))
+
+
+def reset_cb(ctx, reset_type):
+    return 0
+
+
+@vfu_reset_cb_t
+def _reset_cb(ctx, reset_type):
+    return reset_cb(ctx, reset_type)
+
+
+def vfu_setup_device_reset_cb(ctx, cb=_reset_cb):
+    assert ctx is not None
+    return lib.vfu_setup_device_reset_cb(ctx, c.cast(cb, vfu_reset_cb_t))
+
+
+def prepare_ctx_for_dma(dma_register=__dma_register,
+                        dma_unregister=__dma_unregister, quiesce=_quiesce_cb,
+                        reset=_reset_cb):
     ctx = vfu_create_ctx(flags=LIBVFIO_USER_FLAG_ATTACH_NB)
     assert ctx is not None
 
@@ -754,6 +866,23 @@ def prepare_ctx_for_dma():
     ret = vfu_setup_device_dma(ctx, dma_register, dma_unregister)
     assert ret == 0
 
+    if quiesce is not None:
+        vfu_setup_device_quiesce_cb(ctx, quiesce)
+
+    if reset is not None:
+        ret = vfu_setup_device_reset_cb(ctx, reset)
+        assert ret == 0
+
+    f = tempfile.TemporaryFile()
+    f.truncate(0x2000)
+
+    mmap_areas = [(0x1000, 0x1000)]
+
+    ret = vfu_setup_region(ctx, index=VFU_PCI_DEV_MIGR_REGION_IDX, size=0x2000,
+                           flags=VFU_REGION_FLAG_RW, mmap_areas=mmap_areas,
+                           fd=f.fileno())
+    assert ret == 0
+
     ret = vfu_realize_ctx(ctx)
     assert ret == 0
 
@@ -769,7 +898,15 @@ msg_id = 1
 
 @c.CFUNCTYPE(None, c.c_void_p, c.c_int, c.c_char_p)
 def log(ctx, level, msg):
-    print(msg.decode("utf-8"))
+    lvl2str = {syslog.LOG_EMERG: "EMERGENCY",
+                syslog.LOG_ALERT: "ALERT",
+                syslog.LOG_CRIT: "CRITICAL",
+                syslog.LOG_ERR: "ERROR",
+                syslog.LOG_WARNING: "WANRING",
+                syslog.LOG_NOTICE: "NOTICE",
+                syslog.LOG_INFO: "INFO",
+                syslog.LOG_DEBUG: "DEBUG"}
+    print(lvl2str[level] + ": " + msg.decode("utf-8"))
 
 
 def vfio_user_header(cmd, size, no_reply=False, error=False, error_no=0):
@@ -803,15 +940,21 @@ def vfu_realize_ctx(ctx):
 def vfu_attach_ctx(ctx, expect=0):
     ret = lib.vfu_attach_ctx(ctx)
     if expect == 0:
-        assert ret == 0
+        assert ret == 0, "failed to attach: %s" % os.strerror(c.get_errno())
     else:
         assert ret == -1
         assert c.get_errno() == expect
     return ret
 
 
-def vfu_run_ctx(ctx):
-    return lib.vfu_run_ctx(ctx)
+def vfu_run_ctx(ctx, expect_errno=None):
+    ret = lib.vfu_run_ctx(ctx)
+    if expect_errno is not None:
+        assert ret < 0 and expect_errno == c.get_errno(), \
+            "expected '%s' (%d), actual '%s' (%s)" % \
+                (os.strerror(expect_errno), expect_errno,
+                 os.strerror(c.get_errno()), c.get_errno())
+    return ret
 
 
 def vfu_destroy_ctx(ctx):
@@ -821,7 +964,16 @@ def vfu_destroy_ctx(ctx):
         os.remove(SOCK_PATH)
 
 
-def vfu_setup_region(ctx, index, size, cb=None, flags=0,
+def pci_region_cb(ctx, buf, count, offset, is_write):
+    pass
+
+
+@vfu_region_access_cb_t
+def __pci_region_cb(ctx, buf, count, offset, is_write):
+    return pci_region_cb(ctx, buf, count, offset, is_write)
+
+
+def vfu_setup_region(ctx, index, size, cb=__pci_region_cb, flags=0,
                      mmap_areas=None, nr_mmap_areas=None, fd=-1, offset=0):
     assert ctx is not None
 
@@ -851,11 +1003,6 @@ def vfu_setup_region(ctx, index, size, cb=None, flags=0,
     return ret
 
 
-def vfu_setup_device_reset_cb(ctx, cb):
-    assert ctx is not None
-    return lib.vfu_setup_device_reset_cb(ctx, c.cast(cb, vfu_reset_cb_t))
-
-
 def vfu_setup_device_nr_irqs(ctx, irqtype, count):
     assert ctx is not None
     return lib.vfu_setup_device_nr_irqs(ctx, irqtype, count)
@@ -901,22 +1048,75 @@ def vfu_setup_device_dma(ctx, register_cb=None, unregister_cb=None):
                                                 vfu_dma_unregister_cb_t))
 
 
+# FIXME some of the migration arguments are probably wrong as in the C version
+# they're pointer. Check how we handle the read/write region callbacks.
+
+def migr_trans_cb(ctx, state):
+    pass
+
+
+@transition_cb_t
+def __migr_trans_cb(ctx, state):
+    return migr_trans_cb(ctx, state)
+
+
+def migr_get_pending_bytes_cb(ctx):
+    pass
+
+
+@get_pending_bytes_cb_t
+def __migr_get_pending_bytes_cb(ctx):
+    return migr_get_pending_bytes_cb(ctx)
+
+
+def migr_prepare_data_cb(ctx, offset, size):
+    pass
+
+
+@prepare_data_cb_t
+def __migr_prepare_data_cb(ctx, offset, size):
+    return migr_prepare_data_cb(ctx, offset, size)
+
+
+def migr_read_data_cb(ctx, buf, count, offset):
+    pass
+
+
+@read_data_cb_t
+def __migr_read_data_cb(ctx, buf, count, offset):
+    return migr_read_data_cb(ctx, buf, count, offset)
+
+
+def migr_write_data_cb(ctx, buf, count, offset):
+    pass
+
+
+@write_data_cb_t
+def __migr_write_data_cb(ctx, buf, count, offset):
+    return migr_write_data_cb(ctx, buf, count, offset)
+
+
+def migr_data_written_cb(ctx, count):
+    pass
+
+
+@data_written_cb_t
+def __migr_data_written_cb(ctx, count):
+    return migr_data_written_cb(ctx, count)
+
+
 def vfu_setup_device_migration_callbacks(ctx, cbs=None, offset=0):
     assert ctx is not None
 
-    @c.CFUNCTYPE(c.c_int)
-    def stub():
-        return 0
-
     if not cbs:
         cbs = vfu_migration_callbacks_t()
         cbs.version = VFU_MIGR_CALLBACKS_VERS
-        cbs.transition = c.cast(stub, transition_cb_t)
-        cbs.get_pending_bytes = c.cast(stub, get_pending_bytes_cb_t)
-        cbs.prepare_data = c.cast(stub, prepare_data_cb_t)
-        cbs.read_data = c.cast(stub, read_data_cb_t)
-        cbs.write_data = c.cast(stub, write_data_cb_t)
-        cbs.data_written = c.cast(stub, data_written_cb_t)
+        cbs.transition = __migr_trans_cb
+        cbs.get_pending_bytes = __migr_get_pending_bytes_cb
+        cbs.prepare_data = __migr_prepare_data_cb
+        cbs.read_data = __migr_read_data_cb
+        cbs.write_data = __migr_write_data_cb
+        cbs.data_written = __migr_data_written_cb
 
     return lib.vfu_setup_device_migration_callbacks(ctx, cbs, offset)
 
@@ -945,7 +1145,15 @@ def vfu_create_ioeventfd(ctx, region_idx, fd, offset, size, flags, datamatch):
                                     flags, datamatch)
 
 
-def vfu_migr_done(ctx, err):
-    return lib.vfu_migr_done(ctx, err)
+def vfu_device_quiesced(ctx, err):
+    return lib.vfu_device_quiesced(ctx, err)
+
+
+def fail_with_errno(err):
+    def side_effect(args, *kwargs):
+        c.set_errno(err)
+        return -1
+    return side_effect
+
 
 # ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_destroy.py b/test/py/test_destroy.py
new file mode 100644
index 0000000..713c1fb
--- /dev/null
+++ b/test/py/test_destroy.py
@@ -0,0 +1,58 @@
+#
+# Copyright (c) 2021 Nutanix Inc. All rights reserved.
+#
+# Authors: Thanos Makatos <thanos.makatos@nutanix.com>
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#      * Redistributions of source code must retain the above copyright
+#        notice, this list of conditions and the following disclaimer.
+#      * Redistributions in binary form must reproduce the above copyright
+#        notice, this list of conditions and the following disclaimer in the
+#        documentation and/or other materials provided with the distribution.
+#      * Neither the name of Nutanix nor the names of its contributors may be
+#        used to endorse or promote products derived from this software without
+#        specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
+#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#  SERVICESLOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+#  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+#  OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+#  DAMAGE.
+#
+
+from libvfio_user import *
+from unittest.mock import patch
+
+
+ctx = None
+
+
+def setup_function(function):
+    global ctx, sock
+    ctx = prepare_ctx_for_dma()
+    assert ctx is not None
+    sock = connect_client(ctx)
+
+
+def teardown_function(function):
+    pass
+
+
+@patch('libvfio_user.quiesce_cb')
+@patch('libvfio_user.reset_cb', return_value=0)
+def test_destroy_ctx(mock_reset, mock_quiesce):
+    """Checks that destroying a context doesn't call the quiesce callback."""
+
+    vfu_destroy_ctx(ctx)
+    assert mock_quiesce.call_count == 0
+    mock_reset.assert_called_once_with(ctx, VFU_RESET_LOST_CONN)
+
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_device_get_info.py b/test/py/test_device_get_info.py
index c41d382..ae44555 100644
--- a/test/py/test_device_get_info.py
+++ b/test/py/test_device_get_info.py
@@ -53,14 +53,16 @@ def test_device_get_info():
 
     payload = struct.pack("II", 0, 0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_GET_INFO, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_GET_INFO, payload,
+        expect_reply_errno=errno.EINVAL)
 
     # bad argsz
 
     payload = vfio_user_device_info(argsz=8, flags=0,
                                     num_regions=0, num_irqs=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_GET_INFO, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_GET_INFO, payload,
+        expect_reply_errno=errno.EINVAL)
 
     # valid with larger argsz
 
@@ -79,3 +81,5 @@ def test_device_get_info():
     disconnect_client(ctx, sock)
 
     vfu_destroy_ctx(ctx)
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_device_get_irq_info.py b/test/py/test_device_get_irq_info.py
index c45ad7f..fd00544 100644
--- a/test/py/test_device_get_irq_info.py
+++ b/test/py/test_device_get_irq_info.py
@@ -61,19 +61,22 @@ def test_device_get_irq_info_setup():
 def test_device_get_irq_info_bad_in():
     payload = struct.pack("II", 0, 0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload,
+        expect_reply_errno=errno.EINVAL)
 
     # bad argsz
     payload = vfio_irq_info(argsz=8, flags=0, index=VFU_DEV_REQ_IRQ,
                             count=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload,
+        expect_reply_errno=errno.EINVAL)
 
     # bad index
     payload = vfio_irq_info(argsz=argsz, flags=0, index=VFU_DEV_NUM_IRQS,
                             count=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_get_irq_info():
@@ -126,3 +129,5 @@ def test_device_get_irq_info_cleanup():
     disconnect_client(ctx, sock)
 
     vfu_destroy_ctx(ctx)
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_device_get_region_info.py b/test/py/test_device_get_region_info.py
index a99f42c..d00935b 100644
--- a/test/py/test_device_get_region_info.py
+++ b/test/py/test_device_get_region_info.py
@@ -97,7 +97,7 @@ def test_device_get_region_info_short_write():
     payload = struct.pack("II", 0, 0)
 
     msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_INFO, payload,
-        expect=errno.EINVAL)
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_get_region_info_bad_argsz():
@@ -107,7 +107,7 @@ def test_device_get_region_info_bad_argsz():
                                size=0, offset=0)
 
     msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_INFO, payload,
-        expect=errno.EINVAL)
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_get_region_info_bad_index():
@@ -117,7 +117,7 @@ def test_device_get_region_info_bad_index():
                                size=0, offset=0)
 
     msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_INFO, payload,
-        expect=errno.EINVAL)
+        expect_reply_errno=errno.EINVAL)
 
 
 # python tests use max client fds of 8, but this region has 9 mmap areas.
@@ -128,7 +128,7 @@ def test_device_get_region_info_caps_too_few_fds():
     payload = bytes(payload) + b'\0' * (192 - 32)
 
     msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_INFO, payload,
-        expect=errno.ENOSPC)
+        expect_reply_errno=errno.ENOSPC)
 
 
 def test_device_get_region_info_larger_argsz():
@@ -257,3 +257,5 @@ def test_device_get_region_info_migr():
 
 def test_device_get_region_info_cleanup():
     vfu_destroy_ctx(ctx)
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_device_get_region_info_zero_size.py b/test/py/test_device_get_region_info_zero_size.py
index d80152c..3defff1 100644
--- a/test/py/test_device_get_region_info_zero_size.py
+++ b/test/py/test_device_get_region_info_zero_size.py
@@ -75,3 +75,5 @@ def test_device_get_region_info_zero_sized_region():
         assert info.offset == 0
 
     vfu_destroy_ctx(ctx)
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_device_get_region_io_fds.py b/test/py/test_device_get_region_io_fds.py
index bffe508..9e2a09d 100644
--- a/test/py/test_device_get_region_io_fds.py
+++ b/test/py/test_device_get_region_io_fds.py
@@ -28,12 +28,10 @@
 #
 
 from libvfio_user import *
-import ctypes as c
 import errno
 import tempfile
 import os
 import struct
-import ctypes
 
 ctx = None
 sock = None
@@ -90,7 +88,7 @@ def test_device_get_region_io_fds_bad_flags():
                 index=VFU_PCI_DEV_BAR2_REGION_IDX, count=0)
 
     msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_IO_FDS, payload,
-        expect=errno.EINVAL)
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_get_region_io_fds_bad_count():
@@ -101,7 +99,7 @@ def test_device_get_region_io_fds_bad_count():
                 index=VFU_PCI_DEV_BAR2_REGION_IDX, count=1)
 
     msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_IO_FDS, payload,
-        expect=errno.EINVAL)
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_get_region_io_fds_buffer_too_small():
@@ -111,7 +109,7 @@ def test_device_get_region_io_fds_buffer_too_small():
             index=VFU_PCI_DEV_BAR2_REGION_IDX, count=1)
 
     msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_IO_FDS, payload,
-        expect=errno.EINVAL)
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_get_region_io_fds_buffer_too_large():
@@ -122,7 +120,7 @@ def test_device_get_region_io_fds_buffer_too_large():
                                             count=1)
 
     msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_IO_FDS, payload,
-        expect=errno.EINVAL)
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_get_region_io_fds_no_fds():
@@ -130,7 +128,8 @@ def test_device_get_region_io_fds_no_fds():
     payload = vfio_user_region_io_fds_request(argsz=512, flags=0,
                                 index=VFU_PCI_DEV_BAR1_REGION_IDX, count=0)
 
-    ret = msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_IO_FDS, payload, expect=0)
+    ret = msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_IO_FDS, payload,
+              expect_reply_errno=0)
 
     reply, ret = vfio_user_region_io_fds_reply.pop_from_buffer(ret)
 
@@ -146,7 +145,7 @@ def test_device_get_region_io_fds_no_regions_setup():
                                 index=VFU_PCI_DEV_BAR3_REGION_IDX, count=0)
 
     msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_IO_FDS, payload,
-        expect=errno.EINVAL)
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_get_region_io_fds_region_no_mmap():
@@ -155,7 +154,7 @@ def test_device_get_region_io_fds_region_no_mmap():
                                 index=VFU_PCI_DEV_BAR5_REGION_IDX, count=0)
 
     ret = msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_IO_FDS, payload,
-              expect=0)
+              expect_reply_errno=0)
 
     reply, ret = vfio_user_region_io_fds_reply.pop_from_buffer(ret)
 
@@ -171,7 +170,7 @@ def test_device_get_region_io_fds_region_out_of_range():
                                               index=512, count=0)
 
     msg(ctx, sock, VFIO_USER_DEVICE_GET_REGION_IO_FDS, payload,
-        expect=errno.EINVAL)
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_get_region_io_fds_fds_read_write():
@@ -325,3 +324,5 @@ def test_device_get_region_info_cleanup():
     for i in fds:
         os.close(i)
     vfu_destroy_ctx(ctx)
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_device_set_irqs.py b/test/py/test_device_set_irqs.py
index 6d47778..81239a8 100644
--- a/test/py/test_device_set_irqs.py
+++ b/test/py/test_device_set_irqs.py
@@ -69,7 +69,8 @@ def test_device_set_irqs_no_irq_set():
 def test_device_set_irqs_short_write():
     payload = struct.pack("II", 0, 0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_argsz():
@@ -77,7 +78,8 @@ def test_device_set_irqs_bad_argsz():
                            VFIO_IRQ_SET_DATA_NONE, index=VFU_DEV_REQ_IRQ,
                            start=0, count=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_index():
@@ -85,7 +87,8 @@ def test_device_set_irqs_bad_index():
                            VFIO_IRQ_SET_DATA_NONE, index=VFU_DEV_NUM_IRQS,
                            start=0, count=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_flags_MASK_and_UNMASK():
@@ -93,7 +96,8 @@ def test_device_set_irqs_bad_flags_MASK_and_UNMASK():
                            VFIO_IRQ_SET_ACTION_UNMASK, index=VFU_DEV_MSIX_IRQ,
                            start=0, count=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_flags_DATA_NONE_and_DATA_BOOL():
@@ -101,7 +105,8 @@ def test_device_set_irqs_bad_flags_DATA_NONE_and_DATA_BOOL():
                            VFIO_IRQ_SET_DATA_NONE | VFIO_IRQ_SET_DATA_BOOL,
                            index=VFU_DEV_MSIX_IRQ, start=0, count=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_start_count_range():
@@ -109,7 +114,8 @@ def test_device_set_irqs_bad_start_count_range():
                            VFIO_IRQ_SET_DATA_NONE, index=VFU_DEV_MSIX_IRQ,
                            start=2047, count=2)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_start_count_range2():
@@ -117,7 +123,8 @@ def test_device_set_irqs_bad_start_count_range2():
                            VFIO_IRQ_SET_DATA_NONE, index=VFU_DEV_MSIX_IRQ,
                            start=2049, count=1)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_action_for_err_irq():
@@ -125,7 +132,8 @@ def test_device_set_irqs_bad_action_for_err_irq():
                            VFIO_IRQ_SET_DATA_NONE, index=VFU_DEV_ERR_IRQ,
                            start=0, count=1)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_action_for_req_irq():
@@ -133,7 +141,8 @@ def test_device_set_irqs_bad_action_for_req_irq():
                            VFIO_IRQ_SET_DATA_NONE, index=VFU_DEV_REQ_IRQ,
                            start=0, count=1)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_start_for_count_0():
@@ -141,7 +150,8 @@ def test_device_set_irqs_bad_start_for_count_0():
                            VFIO_IRQ_SET_DATA_NONE, index=VFU_DEV_MSIX_IRQ,
                            start=1, count=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_action_for_count_0():
@@ -149,7 +159,8 @@ def test_device_set_irqs_bad_action_for_count_0():
                            VFIO_IRQ_SET_DATA_NONE, index=VFU_DEV_MSIX_IRQ,
                            start=0, count=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_action_and_data_type_for_count_0():
@@ -157,7 +168,8 @@ def test_device_set_irqs_bad_action_and_data_type_for_count_0():
                            VFIO_IRQ_SET_DATA_BOOL, index=VFU_DEV_MSIX_IRQ,
                            start=0, count=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_bad_fds_for_DATA_BOOL():
@@ -169,8 +181,8 @@ def test_device_set_irqs_bad_fds_for_DATA_BOOL():
 
     fd = eventfd()
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL,
-        fds=[fd])
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL, fds=[fd])
 
     os.close(fd)
 
@@ -182,8 +194,8 @@ def test_device_set_irqs_bad_fds_for_DATA_NONE():
 
     fd = eventfd()
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL,
-        fds=[fd])
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL, fds=[fd])
 
     os.close(fd)
 
@@ -195,8 +207,8 @@ def test_device_set_irqs_bad_fds_for_count_2():
 
     fd = eventfd()
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL,
-        fds=[fd])
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL, fds=[fd])
 
     os.close(fd)
 
@@ -231,7 +243,8 @@ def test_device_set_irqs_trigger_bool_too_small():
                            start=0, count=2)
     payload = bytes(payload) + struct.pack("?", False)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_trigger_bool_too_large():
@@ -240,7 +253,8 @@ def test_device_set_irqs_trigger_bool_too_large():
                            start=0, count=2)
     payload = bytes(payload) + struct.pack("???", False, False, False)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DEVICE_SET_IRQS, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_device_set_irqs_enable_update():
@@ -296,3 +310,5 @@ def test_device_set_irqs_enable_trigger_bool():
 
 def test_device_set_irqs_cleanup():
     vfu_destroy_ctx(ctx)
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_dirty_pages.py b/test/py/test_dirty_pages.py
index 6d94a5f..7a84124 100644
--- a/test/py/test_dirty_pages.py
+++ b/test/py/test_dirty_pages.py
@@ -37,12 +37,11 @@ ctx = None
 
 @vfu_dma_register_cb_t
 def dma_register(ctx, info):
-    pass
+    return 0
 
 
 @vfu_dma_unregister_cb_t
 def dma_unregister(ctx, info):
-    pass
     return 0
 
 
@@ -92,21 +91,24 @@ def test_dirty_pages_setup():
 def test_dirty_pages_short_write():
     payload = struct.pack("I", 8)
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dirty_pages_bad_argsz():
     payload = vfio_user_dirty_pages(argsz=4,
         flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_START)
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dirty_pages_start_no_migration():
     payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
         flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_START)
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.ENOTSUP)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.ENOTSUP)
 
 
 def test_dirty_pages_start_bad_flags():
@@ -120,13 +122,15 @@ def test_dirty_pages_start_bad_flags():
         flags=(VFIO_IOMMU_DIRTY_PAGES_FLAG_START |
                VFIO_IOMMU_DIRTY_PAGES_FLAG_STOP))
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.EINVAL)
 
     payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
         flags=(VFIO_IOMMU_DIRTY_PAGES_FLAG_START |
                VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP))
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def start_logging():
@@ -146,7 +150,8 @@ def test_dirty_pages_get_short_read():
     payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
         flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_GET_BITMAP)
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 #
@@ -161,7 +166,8 @@ def test_dirty_pages_get_sub_range():
 
     payload = bytes(dirty_pages) + bytes(br)
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.ENOTSUP)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.ENOTSUP)
 
 
 def test_dirty_pages_get_bad_page_size():
@@ -173,7 +179,8 @@ def test_dirty_pages_get_bad_page_size():
 
     payload = bytes(dirty_pages) + bytes(br)
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dirty_pages_get_bad_bitmap_size():
@@ -185,7 +192,8 @@ def test_dirty_pages_get_bad_bitmap_size():
 
     payload = bytes(dirty_pages) + bytes(br)
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dirty_pages_get_bad_argsz():
@@ -197,7 +205,8 @@ def test_dirty_pages_get_bad_argsz():
 
     payload = bytes(dirty_pages) + bytes(br)
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dirty_pages_get_short_reply():
@@ -230,7 +239,8 @@ def test_get_dirty_page_bitmap_unmapped():
 
     payload = bytes(dirty_pages) + bytes(br)
 
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dirty_pages_get_unmodified():
@@ -353,5 +363,4 @@ def test_dirty_pages_cleanup():
     disconnect_client(ctx, sock)
     vfu_destroy_ctx(ctx)
 
-
 # ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
diff --git a/test/py/test_dma_map.py b/test/py/test_dma_map.py
index f446efd..2a9ac96 100644
--- a/test/py/test_dma_map.py
+++ b/test/py/test_dma_map.py
@@ -27,6 +27,10 @@
 #  DAMAGE.
 #
 
+from unittest import mock
+from unittest.mock import patch
+import mmap
+
 from libvfio_user import *
 import errno
 
@@ -37,26 +41,32 @@ import errno
 ctx = None
 
 
-def test_dma_region_too_big():
-    global ctx
-
+def setup_function(function):
+    global ctx, sock
     ctx = prepare_ctx_for_dma()
     assert ctx is not None
-
     sock = connect_client(ctx)
 
+
+def teardown_function(function):
+    global ctx, sock
+    disconnect_client(ctx, sock)
+    vfu_destroy_ctx(ctx)
+
+
+def test_dma_region_too_big():
+    global ctx, sock
+
     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=0x10000, size=MAX_DMA_SIZE + 4096)
 
-    msg(ctx, sock, VFIO_USER_DMA_MAP, payload, expect=errno.ENOSPC)
-
-    disconnect_client(ctx, sock)
+    msg(ctx, sock, VFIO_USER_DMA_MAP, payload, expect_reply_errno=errno.ENOSPC)
 
 
 def test_dma_region_too_many():
-    sock = connect_client(ctx)
+    global ctx, sock
 
     for i in range(1, MAX_DMA_REGIONS + 2):
         payload = vfio_user_dma_map(argsz=len(vfio_user_dma_map()),
@@ -69,10 +79,154 @@ def test_dma_region_too_many():
         else:
             expect = 0
 
-        msg(ctx, sock, VFIO_USER_DMA_MAP, payload, expect=expect)
+        msg(ctx, sock, VFIO_USER_DMA_MAP, payload, expect_reply_errno=expect)
 
-    disconnect_client(ctx, sock)
 
+@patch('libvfio_user.quiesce_cb', side_effect=fail_with_errno(errno.EBUSY))
+@patch('libvfio_user.dma_register')
+def test_dma_map_busy(mock_dma_register, mock_quiesce):
+    """
+    Checks that during a DMA map operation where the device is initially busy
+    quiescing, and then eventually quiesces, the DMA map operation succeeds.
+    """
 
-def test_dma_region_cleanup():
-    vfu_destroy_ctx(ctx)
+    global ctx, sock
+
+    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=0x10000, size=0x1000)
+
+    msg(ctx, sock, VFIO_USER_DMA_MAP, payload, rsp=False,
+        expect_run_ctx_errno=errno.EBUSY)
+
+    assert mock_dma_register.call_count == 0
+
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
+
+    # check that DMA register callback got called
+    dma_info = vfu_dma_info_t(iovec_t(iov_base=0x10000, iov_len=0x1000),
+        None, iovec_t(None, 0), 0x1000, mmap.PROT_READ | mmap.PROT_WRITE)
+    mock_dma_register.assert_called_once_with(ctx, dma_info)
+
+    get_reply(sock)
+
+    ret = vfu_run_ctx(ctx)
+    assert ret == 0
+
+    # the callback shouldn't be called again
+    mock_dma_register.assert_called_once()
+
+    # check that the DMA region has been added
+    count, sgs = vfu_addr_to_sg(ctx, 0x10000, 0x1000)
+    assert len(sgs) == 1
+    sg = sgs[0]
+    assert sg.dma_addr == 0x10000 and sg.region == 0 and sg.length == 0x1000 \
+        and sg.offset == 0 and sg.writeable
+
+
+# FIXME better move this test and the following to test_request_errors
+
+
+# FIXME need the same test for (1) DMA unmap, (2) device reset, and
+# (3) migration, where quiesce returns EBUSY but replying fails.
+@patch('libvfio_user.reset_cb')
+@patch('libvfio_user.quiesce_cb', return_value=0)
+@patch('libvfio_user.dma_register')
+def test_dma_map_reply_fail(mock_dma_register, mock_quiesce, mock_reset):
+    """Tests mapping a DMA region where the quiesce callback returns 0 and
+    replying fails."""
+
+    global ctx, sock
+
+    # The only chance we have to allow the message to be received but for the
+    # reply to fail is in the DMA map callback, where the message has been
+    # received but reply hasn't been sent yet.
+    def side_effect(ctx, info):
+        sock.close()
+
+    mock_dma_register.side_effect = side_effect
+
+    # Send a DMA map command.
+    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=0x10000, size=0x1000)
+
+    msg(ctx, sock, VFIO_USER_DMA_MAP, payload, rsp=False)
+
+    vfu_run_ctx(ctx, errno.ENOTCONN)
+
+    # TODO not sure whether the following is worth it?
+    try:
+        get_reply(sock)
+    except OSError as e:
+        assert e.errno == errno.EBADF
+    else:
+        assert False
+
+    # 1st call is for adding the DMA region, 2nd call is for the reset
+    mock_quiesce.assert_has_calls([mock.call(ctx)] * 2)
+    mock_reset.assert_called_once_with(ctx, VFU_RESET_LOST_CONN)
+
+    # no need to check that DMA region wasn't added as the context is reset
+
+
+# FIXME need the same test for (1) DMA unmap, (2) device reset, and
+# (3) migration, where quiesce returns EBUSY but replying fails.
+@patch('libvfio_user.reset_cb')
+@patch('libvfio_user.quiesce_cb', side_effect=fail_with_errno(errno.EBUSY))
+@patch('libvfio_user.dma_register')
+def test_dma_map_busy_reply_fail(mock_dma_register, mock_quiesce, mock_reset):
+    """
+    Tests mapping a DMA region where the quiesce callback returns EBUSY and
+    replying fails.
+    """
+
+    global ctx, sock
+
+    # Send a DMA map command.
+    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=0x10000, size=0x1000)
+
+    msg(ctx, sock, VFIO_USER_DMA_MAP, payload, rsp=False,
+        expect_run_ctx_errno=errno.EBUSY)
+
+    mock_quiesce.assert_called_once_with(ctx)
+
+    # pretend there's a connection failure while the device is still quiescing
+    sock.close()
+
+    mock_dma_register.assert_not_called()
+    mock_reset.assert_not_called()
+
+    # device quiesces
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
+
+    dma_info = vfu_dma_info_t(iovec_t(iov_base=0x10000, iov_len=0x1000),
+        None, iovec_t(None, 0), 0x1000, mmap.PROT_READ | mmap.PROT_WRITE)
+    mock_dma_register.assert_called_once_with(ctx, dma_info)
+
+    # device reset callback should be called (by do_reply)
+    mock_reset.assert_called_once_with(ctx, True)
+
+    vfu_run_ctx(ctx, errno.ENOTCONN)
+
+    # callbacks shouldn't be called further
+    mock_quiesce.assert_called_once()
+    mock_dma_register.assert_called_once()
+    mock_reset.assert_called_once()
+
+    # check that the DMA region was NOT added
+    count, sgs = vfu_addr_to_sg(ctx, 0x10000, 0x1000)
+    assert count == -1
+    assert c.get_errno() == errno.ENOENT
+
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_dma_unmap.py b/test/py/test_dma_unmap.py
index a449990..c464ae1 100644
--- a/test/py/test_dma_unmap.py
+++ b/test/py/test_dma_unmap.py
@@ -29,59 +29,60 @@
 #
 
 import errno
-import tempfile
+from unittest.mock import patch
 from libvfio_user import *
 
 ctx = None
 sock = None
 
 
-def test_dma_unmap_setup():
+def setup_function(function):
     global ctx, sock
-
     ctx = prepare_ctx_for_dma()
     assert ctx is not None
-    f = tempfile.TemporaryFile()
-    f.truncate(0x2000)
-
-    mmap_areas = [(0x1000, 0x1000)]
-
-    ret = vfu_setup_region(ctx, index=VFU_PCI_DEV_MIGR_REGION_IDX, size=0x2000,
-                           flags=VFU_REGION_FLAG_RW, mmap_areas=mmap_areas,
-                           fd=f.fileno())
-    assert ret == 0
 
     ret = vfu_realize_ctx(ctx)
     assert ret == 0
 
     sock = connect_client(ctx)
 
-    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=0x1000, size=4096)
 
-    msg(ctx, sock, VFIO_USER_DMA_MAP, payload)
+def teardown_function(function):
+    global ctx, sock
+    disconnect_client(ctx, sock)
+    vfu_destroy_ctx(ctx)
 
 
-def test_dma_unmap_short_write():
+def setup_dma_regions(dma_regions=[(0x0, 0x1000)]):
+    global ctx, sock
+    for dma_region in dma_regions:
+        payload = struct.pack("II", 0, 0)
+        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=dma_region[0], size=dma_region[1])
+        msg(ctx, sock, VFIO_USER_DMA_MAP, payload)
 
-    payload = struct.pack("II", 0, 0)
 
-    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload, expect=errno.EINVAL)
+def test_dma_unmap_short_write():
+    payload = struct.pack("II", 0, 0)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dma_unmap_bad_argsz():
 
     payload = vfio_user_dma_unmap(argsz=8, flags=0, addr=0x1000, size=4096)
-    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dma_unmap_bad_argsz2():
 
     payload = vfio_user_dma_unmap(argsz=SERVER_MAX_DATA_XFER_SIZE + 8, flags=0,
                                   addr=0x1000, size=4096)
-    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dma_unmap_dirty_bad_argsz():
@@ -92,22 +93,26 @@ def test_dma_unmap_dirty_bad_argsz():
     bitmap = vfio_user_bitmap(pgsize=4096, size=(UINT64_MAX - argsz) + 8)
     payload = bytes(unmap) + bytes(bitmap)
 
-    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dma_unmap_dirty_not_tracking():
 
+    setup_dma_regions([(0x1000, 4096)])
     argsz = len(vfio_user_dma_unmap()) + len(vfio_user_bitmap()) + 8
     unmap = vfio_user_dma_unmap(argsz=argsz,
         flags=VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP, addr=0x1000, size=4096)
     bitmap = vfio_user_bitmap(pgsize=4096, size=8)
     payload = bytes(unmap) + bytes(bitmap) + bytes(8)
 
-    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dma_unmap_dirty_not_mapped():
 
+    setup_dma_regions([(0x1000, 4096)])
     vfu_setup_device_migration_callbacks(ctx, offset=0x1000)
     payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
                                     flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_START)
@@ -120,36 +125,61 @@ def test_dma_unmap_dirty_not_mapped():
     bitmap = vfio_user_bitmap(pgsize=4096, size=8)
     payload = bytes(unmap) + bytes(bitmap) + bytes(8)
 
-    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dma_unmap_invalid_flags():
 
+    setup_dma_regions()
     payload = vfio_user_dma_unmap(argsz=len(vfio_user_dma_unmap()),
                                   flags=0x4, addr=0x1000, size=4096)
-    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dma_unmap():
 
+    setup_dma_regions()
     payload = vfio_user_dma_unmap(argsz=len(vfio_user_dma_unmap()),
-                                  flags=0, addr=0x1000, size=4096)
+                                  flags=0, addr=0x0, size=0x1000)
     msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload)
 
 
-def test_dma_unmap_all():
+def test_dma_unmap_invalid_addr():
 
-    for i in range(0, MAX_DMA_REGIONS):
-        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=0x1000 * i, size=4096)
+    setup_dma_regions()
+    payload = vfio_user_dma_unmap(argsz=len(vfio_user_dma_unmap()),
+                                  addr=0x10000, size=4096)
 
-        msg(ctx, sock, VFIO_USER_DMA_MAP, payload)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload,
+        expect_reply_errno=errno.ENOENT)
+
+
+@patch('libvfio_user.quiesce_cb')
+def test_dma_unmap_async(mock_quiesce):
 
+    setup_dma_regions()
+    mock_quiesce.side_effect = fail_with_errno(errno.EBUSY)
     payload = vfio_user_dma_unmap(argsz=len(vfio_user_dma_unmap()),
-        flags=VFIO_DMA_UNMAP_FLAG_ALL, addr=0, size=0)
+                                  flags=0, addr=0x0, size=0x1000)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload, rsp=False,
+        expect_run_ctx_errno=errno.EBUSY)
+
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
 
+    get_reply(sock)
+
+    ret = vfu_run_ctx(ctx)
+    assert ret == 0
+
+
+def test_dma_unmap_all():
+
+    setup_dma_regions((0x1000*i, 0x1000) for i in range(MAX_DMA_REGIONS))
+    payload = vfio_user_dma_unmap(argsz=len(vfio_user_dma_unmap()),
+        flags=VFIO_DMA_UNMAP_FLAG_ALL, addr=0, size=0)
     msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload)
 
 
@@ -158,7 +188,8 @@ def test_dma_unmap_all_invalid_addr():
     payload = vfio_user_dma_unmap(argsz=len(vfio_user_dma_unmap()),
         flags=VFIO_DMA_UNMAP_FLAG_ALL, addr=0x10000, size=4096)
 
-    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload, expect=errno.EINVAL)
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload,
+        expect_reply_errno=errno.EINVAL)
 
 
 def test_dma_unmap_all_invalid_flags():
@@ -167,11 +198,10 @@ def test_dma_unmap_all_invalid_flags():
         flags=(VFIO_DMA_UNMAP_FLAG_ALL | VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP),
                addr=0, size=0)
 
-    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload, expect=errno.EINVAL)
-
+    msg(ctx, sock, VFIO_USER_DMA_UNMAP, payload,
+        expect_reply_errno=errno.EINVAL)
 
-def test_dma_unmap_cleanup():
-    disconnect_client(ctx, sock)
-    vfu_destroy_ctx(ctx)
+# FIXME need to add unit tests that test errors in get_request_header,
+# do_reply, vfu_dma_transfer
 
-# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_irq_trigger.py b/test/py/test_irq_trigger.py
index 39c75af..18469a4 100644
--- a/test/py/test_irq_trigger.py
+++ b/test/py/test_irq_trigger.py
@@ -85,3 +85,5 @@ def test_irq_trigger():
 
 def test_irq_trigger_cleanup():
     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 aa271f9..d70be5e 100644
--- a/test/py/test_migration.py
+++ b/test/py/test_migration.py
@@ -30,23 +30,13 @@
 from libvfio_user import *
 import ctypes as c
 import errno
+from unittest.mock import patch
 
 ctx = None
+sock = 0
 
-global trans_cb_err
-trans_cb_err = 0
 
-
-@transition_cb_t
-def trans_cb(ctx, state):
-    global trans_cb_err
-    if trans_cb_err != 0:
-        c.set_errno(trans_cb_err)
-        return -1
-    return 0
-
-
-def test_migration_setup():
+def setup_function(function):
     global ctx, sock
 
     ctx = vfu_create_ctx(flags=LIBVFIO_USER_FLAG_ATTACH_NB)
@@ -56,29 +46,50 @@ def test_migration_setup():
                            flags=VFU_REGION_FLAG_RW)
     assert ret == 0
 
-    @c.CFUNCTYPE(c.c_int)
-    def stub():
-        return 0
-
-    cbs = vfu_migration_callbacks_t()
-    cbs.version = VFU_MIGR_CALLBACKS_VERS
-    cbs.transition = trans_cb
-    cbs.get_pending_bytes = c.cast(stub, get_pending_bytes_cb_t)
-    cbs.prepare_data = c.cast(stub, prepare_data_cb_t)
-    cbs.read_data = c.cast(stub, read_data_cb_t)
-    cbs.write_data = c.cast(stub, write_data_cb_t)
-    cbs.data_written = c.cast(stub, data_written_cb_t)
-
-    ret = vfu_setup_device_migration_callbacks(ctx, cbs, offset=0x4000)
+    ret = vfu_setup_device_migration_callbacks(ctx, offset=0x4000)
     assert ret == 0
 
+    vfu_setup_device_quiesce_cb(ctx)
+
     ret = vfu_realize_ctx(ctx)
     assert ret == 0
 
     sock = connect_client(ctx)
 
 
-def test_migration_trans_sync():
+def teardown_function(function):
+    global ctx
+    vfu_destroy_ctx(ctx)
+
+
+@patch('libvfio_user.quiesce_cb')
+@patch('libvfio_user.migr_trans_cb')
+def test_migration_bad_access(mock_trans, mock_quiesce):
+    """
+    Tests that attempting to access the migration state register in an
+    non-aligned manner fails.
+
+    This test is important because we tell whether we need to quiesce by
+    checking for a register-sized access, otherwise we'll change migration
+    state without having quiesced.
+    """
+    global ctx, sock
+
+    data = VFIO_DEVICE_STATE_SAVING.to_bytes(c.sizeof(c.c_int), 'little')
+    write_region(ctx, sock, VFU_PCI_DEV_MIGR_REGION_IDX, offset=0,
+                 count=len(data)-1, data=data, expect=errno.EINVAL)
+
+    mock_trans.assert_not_called()
+
+
+@patch('libvfio_user.quiesce_cb')
+@patch('libvfio_user.migr_trans_cb', return_value=0)
+def test_migration_trans_sync(mock_trans, mock_quiesce):
+    """
+    Tests transitioning to the saving state.
+    """
+
+    global ctx, sock
 
     data = VFIO_DEVICE_STATE_SAVING.to_bytes(c.sizeof(c.c_int), 'little')
     write_region(ctx, sock, VFU_PCI_DEV_MIGR_REGION_IDX, offset=0,
@@ -88,10 +99,13 @@ def test_migration_trans_sync():
     assert ret == 0
 
 
-def test_migration_trans_sync_err():
+@patch('libvfio_user.migr_trans_cb', side_effect=fail_with_errno(errno.EPERM))
+def test_migration_trans_sync_err(mock_trans):
+    """
+    Tests the device returning an error when the migration state is written to.
+    """
 
-    global trans_cb_err
-    trans_cb_err = errno.EPERM
+    global ctx, sock
 
     data = VFIO_DEVICE_STATE_SAVING.to_bytes(c.sizeof(c.c_int), 'little')
     write_region(ctx, sock, VFU_PCI_DEV_MIGR_REGION_IDX, offset=0,
@@ -101,20 +115,24 @@ def test_migration_trans_sync_err():
     assert ret == 0
 
 
-def test_migration_trans_async():
+@patch('libvfio_user.quiesce_cb', side_effect=fail_with_errno(errno.EBUSY))
+@patch('libvfio_user.migr_trans_cb', return_value=0)
+def test_migration_trans_async(mock_trans, mock_quiesce):
+    """
+    Tests transitioning to the saving state where the device is initially busy
+    quiescing.
+    """
 
-    global trans_cb_err
-    trans_cb_err = errno.EBUSY
+    global ctx, sock
+    mock_quiesce
 
     data = VFIO_DEVICE_STATE_SAVING.to_bytes(c.sizeof(c.c_int), 'little')
     write_region(ctx, sock, VFU_PCI_DEV_MIGR_REGION_IDX, offset=0,
-                 count=len(data), data=data, rsp=False)
-
-    ret = vfu_run_ctx(ctx)
-    assert ret == -1
-    assert c.get_errno() == errno.EBUSY
+                 count=len(data), data=data, rsp=False,
+                 expect_run_ctx_errno=errno.EBUSY)
 
-    vfu_migr_done(ctx, 0)
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
 
     get_reply(sock)
 
@@ -122,23 +140,27 @@ def test_migration_trans_async():
     assert ret == 0
 
 
-def test_migration_trans_async_err():
+@patch('libvfio_user.quiesce_cb', side_effect=fail_with_errno(errno.EBUSY))
+@patch('libvfio_user.migr_trans_cb', side_effect=fail_with_errno(errno.ENOTTY))
+def test_migration_trans_async_err(mock_trans, mock_quiesce):
+    """
+    Tests writing to the migration state register, the device not being able to
+    immediately quiesce, and then finally the device failing to transition to
+    the new migration state.
+    """
 
-    global trans_cb_err
-    trans_cb_err = errno.EBUSY
+    global ctx, sock
 
     data = VFIO_DEVICE_STATE_RUNNING.to_bytes(c.sizeof(c.c_int), 'little')
     write_region(ctx, sock, VFU_PCI_DEV_MIGR_REGION_IDX, offset=0,
-                 count=len(data), data=data, rsp=False)
-
-    ret = vfu_run_ctx(ctx)
-    assert ret == -1
-    assert c.get_errno() == errno.EBUSY
+                 count=len(data), data=data, rsp=False,
+                 expect_run_ctx_errno=errno.EBUSY)
 
-    vfu_migr_done(ctx, errno.ENOTTY)
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
 
+    print("waiting for reply")
     get_reply(sock, errno.ENOTTY)
-
-    vfu_destroy_ctx(ctx)
+    print("received reply")
 
 # ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_negotiate.py b/test/py/test_negotiate.py
index 0541d0e..348bf68 100644
--- a/test/py/test_negotiate.py
+++ b/test/py/test_negotiate.py
@@ -201,3 +201,5 @@ def test_valid_negotiate_json():
 
 def test_destroying():
     vfu_destroy_ctx(ctx)
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
diff --git a/test/py/test_pci_caps.py b/test/py/test_pci_caps.py
index 5267246..effd6d3 100644
--- a/test/py/test_pci_caps.py
+++ b/test/py/test_pci_caps.py
@@ -27,6 +27,7 @@
 #  DAMAGE.
 #
 
+from unittest.mock import patch
 from libvfio_user import *
 import ctypes as c
 import errno
@@ -34,21 +35,37 @@ import errno
 ctx = None
 
 
-def test_pci_cap_setup():
+def setup_function(function):
     global ctx
 
     ctx = vfu_create_ctx(flags=LIBVFIO_USER_FLAG_ATTACH_NB)
     assert ctx is not None
-
-    ret = vfu_pci_init(ctx, pci_type=VFU_PCI_TYPE_CONVENTIONAL)
+    ret = vfu_setup_device_reset_cb(ctx)
     assert ret == 0
+    vfu_setup_device_quiesce_cb(ctx)
 
-    ret = vfu_setup_region(ctx, index=VFU_PCI_DEV_CFG_REGION_IDX,
-                           size=PCI_CFG_SPACE_SIZE, flags=VFU_REGION_FLAG_RW)
+
+def teardown_function(function):
+    vfu_destroy_ctx(ctx)
+
+
+def setup_pci_dev(config_space=True, realize=False):
+    global ctx
+    ret = vfu_pci_init(ctx, pci_type=VFU_PCI_TYPE_CONVENTIONAL)
     assert ret == 0
+    if config_space:
+        ret = vfu_setup_region(ctx, index=VFU_PCI_DEV_CFG_REGION_IDX,
+                               size=PCI_CFG_SPACE_SIZE,
+                               flags=VFU_REGION_FLAG_RW)
+        assert ret == 0
+    if realize:
+        ret = vfu_realize_ctx(ctx)
+        assert ret == 0
 
 
 def test_pci_cap_bad_flags():
+    """Tests adding a PCI capability with bad VFU_CAP_FLAG_ flags."""
+    setup_pci_dev()
     pos = vfu_pci_add_capability(ctx, pos=0, flags=999,
               data=struct.pack("ccHH", to_byte(PCI_CAP_ID_PM), b'\0', 0, 0))
     assert pos == -1
@@ -56,6 +73,10 @@ def test_pci_cap_bad_flags():
 
 
 def test_pci_cap_no_cb():
+    """
+    Tests adding a PCI capability VFU_CAP_FLAG_CALLBACK without a callback.
+    """
+    setup_pci_dev(config_space=False)
     pos = vfu_pci_add_capability(ctx, pos=0, flags=VFU_CAP_FLAG_CALLBACK,
               data=struct.pack("ccHH", to_byte(PCI_CAP_ID_PM), b'\0', 0, 0))
     assert pos == -1
@@ -63,6 +84,8 @@ def test_pci_cap_no_cb():
 
 
 def test_pci_cap_unknown_cap():
+    """Tests adding an unknown PCI capability."""
+    setup_pci_dev()
     pos = vfu_pci_add_capability(ctx, pos=0, flags=0,
               data=struct.pack("ccHH", b'\x81', b'\0', 0, 0))
     assert pos == -1
@@ -70,37 +93,21 @@ def test_pci_cap_unknown_cap():
 
 
 def test_pci_cap_bad_pos():
+    """Tests adding a PCI capability at an invalid position."""
+    setup_pci_dev()
     pos = vfu_pci_add_capability(ctx, pos=PCI_CFG_SPACE_SIZE, flags=0,
               data=struct.pack("ccHH", to_byte(PCI_CAP_ID_PM), b'\0', 0, 0))
     assert pos == -1
     assert c.get_errno() == errno.EINVAL
 
 
-@vfu_region_access_cb_t
-def pci_region_cb(ctx, buf, count, offset, is_write):
+def __pci_region_cb(ctx, buf, count, offset, is_write):
     if not is_write:
         return read_pci_cfg_space(ctx, buf, count, offset)
 
     return write_pci_cfg_space(ctx, buf, count, offset)
 
 
-def test_pci_cap_setup_cb():
-    global ctx
-
-    vfu_destroy_ctx(ctx)
-
-    ctx = vfu_create_ctx(flags=LIBVFIO_USER_FLAG_ATTACH_NB)
-    assert ctx is not None
-
-    ret = vfu_pci_init(ctx, pci_type=VFU_PCI_TYPE_CONVENTIONAL)
-    assert ret == 0
-
-    ret = vfu_setup_region(ctx, index=VFU_PCI_DEV_CFG_REGION_IDX,
-                           size=PCI_CFG_SPACE_SIZE, cb=pci_region_cb,
-                           flags=VFU_REGION_FLAG_RW)
-    assert ret == 0
-
-
 cap_offsets = (
     PCI_STD_HEADER_SIZEOF,
     PCI_STD_HEADER_SIZEOF + PCI_PM_SIZEOF,
@@ -112,7 +119,9 @@ cap_offsets = (
 )
 
 
-def test_add_caps():
+@patch("libvfio_user.pci_region_cb", side_effect=__pci_region_cb)
+def test_add_caps(mock_pci_region_cb):
+    setup_pci_dev()
     pos = vfu_pci_add_capability(ctx, pos=0, flags=0,
               data=struct.pack("ccHH", to_byte(PCI_CAP_ID_PM), b'\0', 0, 0))
     assert pos == cap_offsets[0]
@@ -142,8 +151,17 @@ def test_add_caps():
     ret = vfu_realize_ctx(ctx)
     assert ret == 0
 
+    __test_find_caps()
 
-def test_find_caps():
+    sock = connect_client(ctx)
+
+    __test_pci_cap_write_hdr(sock)
+    __test_pci_cap_readonly(sock)
+    __test_pci_cap_callback(sock)
+    __test_pci_cap_write_pmcs(sock)
+
+
+def __test_find_caps():
     offset = vfu_pci_find_capability(ctx, False, PCI_CAP_ID_PM)
     assert offset == cap_offsets[0]
 
@@ -202,21 +220,15 @@ def test_find_caps():
     assert c.get_errno() == errno.ENOENT
 
 
-def test_pci_cap_write_hdr():
-    sock = connect_client(ctx)
-
+def __test_pci_cap_write_hdr(sock):
     # offset of struct cap_hdr
     offset = cap_offsets[0]
     data = b'\x01'
     write_region(ctx, sock, VFU_PCI_DEV_CFG_REGION_IDX, offset=offset,
                  count=len(data), data=data, expect=errno.EPERM)
 
-    disconnect_client(ctx, sock)
-
-
-def test_pci_cap_readonly():
-    sock = connect_client(ctx)
 
+def __test_pci_cap_readonly(sock):
     # start of vendor payload
     offset = cap_offsets[1] + 2
     data = b'\x01'
@@ -229,12 +241,8 @@ def test_pci_cap_readonly():
                           count=3)
     assert payload == b'abc'
 
-    disconnect_client(ctx, sock)
-
-
-def test_pci_cap_callback():
-    sock = connect_client(ctx)
 
+def __test_pci_cap_callback(sock):
     # offsetof(struct vsc, data)
     offset = cap_offsets[2] + 3
     data = b"Hello world."
@@ -251,11 +259,8 @@ def test_pci_cap_callback():
                           count=len(data))
     assert payload == data
 
-    disconnect_client(ctx, sock)
 
-
-def test_pci_cap_write_pmcs():
-    sock = connect_client(ctx)
+def __test_pci_cap_write_pmcs(sock):
 
     # struct pc
 
@@ -305,40 +310,43 @@ def test_pci_cap_write_pmcs():
     write_region(ctx, sock, VFU_PCI_DEV_CFG_REGION_IDX, offset=offset,
                  count=len(data), data=data, expect=errno.ENOTSUP)
 
-    disconnect_client(ctx, sock)
-
-
-reset_flag = -1
-
-
-@c.CFUNCTYPE(c.c_int, c.c_void_p, c.c_int)
-def vfu_reset_cb(ctx, reset_type):
-    assert reset_type == VFU_RESET_PCI_FLR or reset_type == VFU_RESET_LOST_CONN
-    global reset_flag
-    reset_flag = reset_type
-    return 0
-
-
-def test_pci_cap_write_px():
-    sock = connect_client(ctx)
-    ret = vfu_setup_device_reset_cb(ctx, vfu_reset_cb)
-    assert ret == 0
 
+def _setup_flrc(ctx):
     # flrc
     cap = struct.pack("ccHHcc52c", to_byte(PCI_CAP_ID_EXP), b'\0', 0, 0, b'\0',
                       b'\x10', *[b'\0' for _ in range(52)])
-    pos = vfu_pci_add_capability(ctx, pos=cap_offsets[5], flags=0, data=cap)
-    assert pos == cap_offsets[5]
+    # FIXME adding capability after we've realized the device only works
+    # because of bug #618.
+    pos = vfu_pci_add_capability(ctx, pos=0, flags=0, data=cap)
+    assert pos == PCI_STD_HEADER_SIZEOF
+
+
+@patch("libvfio_user.reset_cb", return_value=0)
+@patch('libvfio_user.quiesce_cb')
+def test_pci_cap_write_px(mock_quiesce, mock_reset):
+    """
+    Tests function level reset.
+    """
+    setup_pci_dev(realize=True)
+    sock = connect_client(ctx)
+
+    _setup_flrc(ctx)
 
     # iflr
-    offset = cap_offsets[5] + 8
+    offset = PCI_STD_HEADER_SIZEOF + 8
     data = b'\x00\x80'
     write_region(ctx, sock, VFU_PCI_DEV_CFG_REGION_IDX, offset=offset,
                  count=len(data), data=data)
-    assert reset_flag == VFU_RESET_PCI_FLR
 
-    disconnect_client(ctx, sock)
-    assert reset_flag == VFU_RESET_LOST_CONN
+    mock_quiesce.assert_called_once_with(ctx)
+    mock_reset.assert_called_once_with(ctx, VFU_RESET_PCI_FLR)
+
+    # bad access
+    for _off in (-1, +1):
+        for _len in (-1, +1):
+            write_region(ctx, sock, VFU_PCI_DEV_CFG_REGION_IDX,
+                         offset=offset+_off, count=len(data)+_len, data=data,
+                         expect=errno.EINVAL)
 
 
 def test_pci_cap_write_msix():
@@ -347,7 +355,12 @@ def test_pci_cap_write_msix():
 
 
 def test_pci_cap_write_pxdc2():
+
+    setup_pci_dev(realize=True)
     sock = connect_client(ctx)
+
+    _setup_flrc(ctx)
+
     offset = (vfu_pci_find_capability(ctx, False, PCI_CAP_ID_EXP) +
               PCI_EXP_DEVCTL2)
     data = b'\xde\xad'
@@ -356,10 +369,13 @@ def test_pci_cap_write_pxdc2():
     payload = read_region(ctx, sock, VFU_PCI_DEV_CFG_REGION_IDX, offset=offset,
                           count=len(data))
     assert payload == data
-    disconnect_client(ctx, sock)
 
 
 def test_pci_cap_write_pxlc2():
+
+    setup_pci_dev(realize=True)
+    _setup_flrc(ctx)
+
     sock = connect_client(ctx)
     offset = (vfu_pci_find_capability(ctx, False, PCI_CAP_ID_EXP) +
               PCI_EXP_LNKCTL2)
@@ -369,10 +385,6 @@ def test_pci_cap_write_pxlc2():
     payload = read_region(ctx, sock, VFU_PCI_DEV_CFG_REGION_IDX, offset=offset,
                           count=len(data))
     assert payload == data
-    disconnect_client(ctx, sock)
-
 
-def test_pci_cap_cleanup():
-    vfu_destroy_ctx(ctx)
 
 # ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_pci_ext_caps.py b/test/py/test_pci_ext_caps.py
index 94eda0e..8fcadf6 100644
--- a/test/py/test_pci_ext_caps.py
+++ b/test/py/test_pci_ext_caps.py
@@ -315,3 +315,5 @@ def test_pci_ext_cap_write_vendor():
 
 def test_pci_ext_cap_cleanup():
     vfu_destroy_ctx(ctx)
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_quiesce.py b/test/py/test_quiesce.py
new file mode 100644
index 0000000..cf41e36
--- /dev/null
+++ b/test/py/test_quiesce.py
@@ -0,0 +1,109 @@
+#
+# Copyright (c) 2021 Nutanix Inc. All rights reserved.
+#
+# Authors: Thanos Makatos <thanos.makatos@nutanix.com>
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#      * Redistributions of source code must retain the above copyright
+#        notice, this list of conditions and the following disclaimer.
+#      * Redistributions in binary form must reproduce the above copyright
+#        notice, this list of conditions and the following disclaimer in the
+#        documentation and/or other materials provided with the distribution.
+#      * Neither the name of Nutanix nor the names of its contributors may be
+#        used to endorse or promote products derived from this software without
+#        specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
+#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#  SERVICESLOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+#  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+#  OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+#  DAMAGE.
+#
+
+from libvfio_user import *
+import errno
+from unittest.mock import patch
+
+
+ctx = None
+
+
+def setup_function(function):
+    global ctx, sock
+    ctx = prepare_ctx_for_dma()
+    assert ctx is not None
+    sock = connect_client(ctx)
+
+
+def teardown_function(function):
+    global ctx
+    vfu_destroy_ctx(ctx)
+
+
+@patch('libvfio_user.quiesce_cb')
+def test_device_quiesced_no_quiesce_requested(mock_quiesce):
+    """
+    Checks that vfu_device_quiesce returns an error if called when there is
+    no pending quiesce operation.
+    """
+
+    global ctx
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == -1
+    assert c.get_errno() == errno.EINVAL
+    assert mock_quiesce.call_count == 0
+
+
+@patch('libvfio_user.quiesce_cb', side_effect=fail_with_errno(errno.ENOTTY))
+def test_device_quiesce_error(mock_quiesce):
+    """
+    Checks that if the device quiesce callback fails then the operation
+    that requested it also fails with the same error.
+    """
+
+    global ctx, sock
+
+    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=0x10000, size=0x1000)
+
+    msg(ctx, sock, VFIO_USER_DMA_MAP, payload, errno.ENOTTY)
+
+
+@patch('libvfio_user.dma_register')
+@patch('libvfio_user.quiesce_cb', side_effect=fail_with_errno(errno.EBUSY))
+def test_device_quiesce_error_after_busy(mock_quiesce, mock_dma_register):
+    """
+    Checks that the device fails to quiesce after it was busy quiescing.
+    """
+
+    global ctx, sock
+
+    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=0x10000, size=0x1000)
+
+    msg(ctx, sock, VFIO_USER_DMA_MAP, payload, rsp=False,
+        expect_run_ctx_errno=errno.EBUSY)
+
+    ret = vfu_device_quiesced(ctx, errno.ENOTTY)
+    assert ret == 0
+
+    mock_dma_register.assert_not_called()
+
+    # check that the DMA region was NOT added
+    count, sgs = vfu_addr_to_sg(ctx, 0x10000, 0x1000)
+    assert count == -1
+    assert c.get_errno() == errno.ENOENT
+
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_request_errors.py b/test/py/test_request_errors.py
index a734bab..6cd50cb 100644
--- a/test/py/test_request_errors.py
+++ b/test/py/test_request_errors.py
@@ -27,17 +27,17 @@
 #  DAMAGE.
 #
 
+from unittest.mock import patch
 from libvfio_user import *
 import errno
 import os
 
 ctx = None
 sock = None
-
 argsz = len(vfio_irq_set())
 
 
-def test_request_errors_setup():
+def setup_function(function):
     global ctx, sock
 
     ctx = vfu_create_ctx(flags=LIBVFIO_USER_FLAG_ATTACH_NB)
@@ -49,12 +49,29 @@ def test_request_errors_setup():
     ret = vfu_setup_device_nr_irqs(ctx, VFU_DEV_MSIX_IRQ, 2048)
     assert ret == 0
 
+    vfu_setup_device_quiesce_cb(ctx)
+
+    ret = vfu_setup_device_reset_cb(ctx)
+    assert ret == 0
+
+    ret = vfu_setup_region(ctx, index=VFU_PCI_DEV_MIGR_REGION_IDX, size=0x2000,
+                           flags=VFU_REGION_FLAG_RW)
+    assert ret == 0
+
+    ret = vfu_setup_device_migration_callbacks(ctx, offset=0x4000)
+    assert ret == 0
+
     ret = vfu_realize_ctx(ctx)
     assert ret == 0
 
     sock = connect_client(ctx)
 
 
+def teardown_function(function):
+    global ctx
+    vfu_destroy_ctx(ctx)
+
+
 def test_too_small():
     # struct vfio_user_header
     hdr = struct.pack("HHIII", 0xbad1, VFIO_USER_DEVICE_SET_IRQS,
@@ -126,5 +143,97 @@ def test_bad_request_closes_fds():
     os.close(fd2)
 
 
-def test_request_errors_cleanup():
-    vfu_destroy_ctx(ctx)
+@patch('libvfio_user.reset_cb')
+@patch('libvfio_user.quiesce_cb', return_value=0)
+def test_disconnected_socket(mock_quiesce, mock_reset):
+    """Tests that calling vfu_run_ctx on a disconnected socket results in
+    resetting the context and returning ENOTCONN."""
+
+    global ctx, sock
+    sock.close()
+
+    vfu_run_ctx(ctx, errno.ENOTCONN)
+
+    # quiesce callback gets called during reset
+    # FIXME how can we ensure that quiesce is called before reset?
+    mock_quiesce.assert_called_with(ctx)
+    mock_reset.assert_called_with(ctx, VFU_RESET_LOST_CONN)
+
+
+@patch('libvfio_user.quiesce_cb', side_effect=fail_with_errno(errno.EBUSY))
+def test_disconnected_socket_quiesce_busy(mock_quiesce):
+    """Tests that calling vfu_run_ctx on a disconnected socket results in
+    resetting the context which returns EBUSY."""
+
+    global ctx, sock
+    sock.close()
+
+    vfu_run_ctx(ctx, errno.EBUSY)
+
+    # quiesce callback must be called during reset
+    mock_quiesce.assert_called_once_with(ctx)
+
+    # device hasn't finished quiescing
+    for _ in range(0, 3):
+        vfu_run_ctx(ctx, errno.EBUSY)
+
+    # device quiesced
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
+
+    vfu_run_ctx(ctx, errno.ENOTCONN)
+
+    # no further calls to the quiesce callback should have been made
+    mock_quiesce.assert_called_once_with(ctx)
+
+
+@patch('libvfio_user.reset_cb')
+@patch('libvfio_user.quiesce_cb', side_effect=fail_with_errno(errno.EBUSY))
+@patch('libvfio_user.migr_get_pending_bytes_cb')
+def test_reply_fail_quiesce_busy(mock_get_pending_bytes, mock_quiesce,
+                                 mock_reset):
+    """Tests failing to reply and the quiesce callback returning EBUSY."""
+
+    global ctx, sock
+
+    def get_pending_bytes_side_effect(ctx):
+        sock.close()
+        return 0
+    mock_get_pending_bytes.side_effect = get_pending_bytes_side_effect
+
+    # read the get_pending_bytes register, it should close the socket causing
+    # the reply to fail
+    read_region(ctx, sock, VFU_PCI_DEV_MIGR_REGION_IDX,
+                vfio_user_migration_info.pending_bytes.offset,
+                vfio_user_migration_info.pending_bytes.size, rsp=False,
+                expect_run_ctx_errno=errno.EBUSY)
+
+    # vfu_run_ctx will try to reset the context and to do that it needs to
+    # quiesce the device first
+    mock_quiesce.assert_called_once_with(ctx)
+
+    # vfu_run_ctx will be returning EBUSY and nothing should have happened
+    # until the device quiesces
+    for _ in range(0, 3):
+        vfu_run_ctx(ctx, errno.EBUSY)
+    mock_quiesce.assert_called_once_with(ctx)
+    mock_reset.assert_not_called()
+
+    ret = vfu_device_quiesced(ctx, 0)
+    assert ret == 0
+
+    # the device quiesced, reset should should happen now
+    mock_quiesce.assert_called_once_with(ctx)
+    mock_reset.assert_called_once_with(ctx, VFU_RESET_LOST_CONN)
+
+    try:
+        get_reply(sock)
+    except OSError as e:
+        assert e.errno == errno.EBADF
+    else:
+        assert False
+
+    vfu_run_ctx(ctx, errno.ENOTCONN)
+
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/py/test_setup_region.py b/test/py/test_setup_region.py
index 76cd1d9..f06d31c 100644
--- a/test/py/test_setup_region.py
+++ b/test/py/test_setup_region.py
@@ -193,5 +193,29 @@ def test_region_offset_overflow():
     disconnect_client(ctx, sock)
 
 
+def test_access_region_zero_count():
+    global ctx
+
+    ret = vfu_setup_region(ctx, index=VFU_PCI_DEV_BAR0_REGION_IDX,
+                           size=0x1000, flags=VFU_REGION_FLAG_RW)
+    assert ret == 0
+
+    ret = vfu_realize_ctx(ctx)
+    assert ret == 0
+
+    sock = connect_client(ctx)
+
+    payload = read_region(ctx, sock, VFU_PCI_DEV_BAR0_REGION_IDX, offset=0,
+                          count=0)
+    assert payload == b''
+
+    write_region(ctx, sock, VFU_PCI_DEV_BAR0_REGION_IDX, offset=0, count=0,
+                 data=payload)
+
+    disconnect_client(ctx, sock)
+
+
 def test_setup_region_cleanup():
     vfu_destroy_ctx(ctx)
+
+# ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: #
diff --git a/test/unit-tests.c b/test/unit-tests.c
index cab3fed..fcaaa2a 100644
--- a/test/unit-tests.c
+++ b/test/unit-tests.c
@@ -231,7 +231,6 @@ test_handle_dma_unmap(void **state UNUSED)
     expect_value(mock_dma_unregister, vfu_ctx, &vfu_ctx);
     expect_check(mock_dma_unregister, info, check_dma_info,
                  &vfu_ctx.dma->regions[0].info);
-    will_return(mock_dma_unregister, 0);
 
     ret = handle_dma_unmap(&vfu_ctx,
                            mkmsg(VFIO_USER_DMA_UNMAP, &dma_unmap,
@@ -286,7 +285,6 @@ test_dma_controller_remove_region_mapped(void **state UNUSED)
     expect_check(mock_dma_unregister, info, check_dma_info,
                  &vfu_ctx.dma->regions[0].info);
     /* FIXME add unit test when dma_unregister fails */
-    will_return(mock_dma_unregister, 0);
     patch("dma_controller_unmap_region");
     expect_value(dma_controller_unmap_region, dma, vfu_ctx.dma);
     expect_value(dma_controller_unmap_region, region, &vfu_ctx.dma->regions[0]);
@@ -306,7 +304,6 @@ test_dma_controller_remove_region_unmapped(void **state UNUSED)
     expect_value(mock_dma_unregister, vfu_ctx, &vfu_ctx);
     expect_check(mock_dma_unregister, info, check_dma_info,
                  &vfu_ctx.dma->regions[0].info);
-    will_return(mock_dma_unregister, 0);
     patch("dma_controller_unmap_region");
     assert_int_equal(0,
         dma_controller_remove_region(vfu_ctx.dma, (void *)0xdeadbeef, 0x100,