From f2dd09649e31540996fa4e9497693d1b27bc88fe Mon Sep 17 00:00:00 2001
From: Thanos Makatos <thanos.makatos@nutanix.com>
Date: Tue, 30 Nov 2021 14:40:18 +0000
Subject: introduce device quiesce callback (#609)

Signed-off-by: Thanos Makatos <thanos.makatos@nutanix.com>
Reviewed-by: John Leon <john.levon@nutanix.com>
---
 README.md                                        |   5 +
 docs/testing.md                                  |   7 +
 include/libvfio-user.h                           |  94 +++++---
 lib/dma.c                                        |  19 +-
 lib/libvfio-user.c                               | 252 ++++++++++++++------
 lib/migration.c                                  |  24 +-
 lib/migration.h                                  |   4 +
 lib/migration_priv.h                             |   3 -
 lib/pci_caps.c                                   |  19 ++
 lib/pci_caps.h                                   |   4 +
 lib/private.h                                    |  31 ++-
 samples/gpio-pci-idio-16.c                       |   3 +-
 samples/server.c                                 |   5 +-
 test/mocks.c                                     |   3 +-
 test/mocks.h                                     |   2 +-
 test/py/libvfio_user.py                          | 290 +++++++++++++++++++----
 test/py/test_destroy.py                          |  58 +++++
 test/py/test_device_get_info.py                  |   8 +-
 test/py/test_device_get_irq_info.py              |  11 +-
 test/py/test_device_get_region_info.py           |  10 +-
 test/py/test_device_get_region_info_zero_size.py |   2 +
 test/py/test_device_get_region_io_fds.py         |  21 +-
 test/py/test_device_set_irqs.py                  |  56 +++--
 test/py/test_dirty_pages.py                      |  37 +--
 test/py/test_dma_map.py                          | 178 +++++++++++++-
 test/py/test_dma_unmap.py                        | 114 +++++----
 test/py/test_irq_trigger.py                      |   2 +
 test/py/test_migration.py                        | 124 ++++++----
 test/py/test_negotiate.py                        |   2 +
 test/py/test_pci_caps.py                         | 156 ++++++------
 test/py/test_pci_ext_caps.py                     |   2 +
 test/py/test_quiesce.py                          | 109 +++++++++
 test/py/test_request_errors.py                   | 117 ++++++++-
 test/py/test_setup_region.py                     |  24 ++
 test/unit-tests.c                                |   3 -
 35 files changed, 1380 insertions(+), 419 deletions(-)
 create mode 100644 test/py/test_destroy.py
 create mode 100644 test/py/test_quiesce.py

diff --git a/README.md b/README.md
index 02d905e..2c7062c 100644
--- a/README.md
+++ b/README.md
@@ -74,6 +74,11 @@ The kernel headers are necessary because VFIO structs and defines are reused.
 
 Finally build your program and link with `libvfio-user.so`.
 
+Coverity
+========
+
+`make coverity` automatically uploads a new coverity build.
+
 Supported features
 ==================
 
diff --git a/docs/testing.md b/docs/testing.md
index 6da018e..71a80b6 100644
--- a/docs/testing.md
+++ b/docs/testing.md
@@ -15,3 +15,10 @@ The master branch is run through [Coverity](scan.coverity.com) when a new PR
 lands.
 
 You can also run `make gcov` to get code coverage reports.
+
+Debugging Test Errors
+---------------------
+
+Sometimes debugging Valgrind errors on Python unit tests can be tricky. To
+run specific tests use the pytest `-k` option in `PYTESTCMD` in the Makefile.
+
diff --git a/include/libvfio-user.h b/include/libvfio-user.h
index b914ccc..a67fe30 100644
--- a/include/libvfio-user.h
+++ b/include/libvfio-user.h
@@ -163,13 +163,16 @@ vfu_get_poll_fd(vfu_ctx_t *vfu_ctx);
  *          with errno set as follows:
  *
  * ENOTCONN: client closed connection, vfu_attach_ctx() should be called again
+ * EBUSY: the device was asked to quiesce and is still quiescing
  * Other errno values are also possible.
  */
 int
 vfu_run_ctx(vfu_ctx_t *vfu_ctx);
 
 /**
- * Destroys libvfio-user context.
+ * Destroys libvfio-user context. During this call the device must already be
+ * in quiesced state; the quiesce callback is not called. Any other device
+ * callback can be called.
  *
  * @vfu_ctx: the libvfio-user context to destroy
  */
@@ -341,6 +344,59 @@ typedef enum vfu_reset_type {
 } vfu_reset_type_t;
 
 /*
+ * Device callback for quiescing the device.
+ *
+ * vfu_run_ctx uses this callback to request from the device to quiesce its
+ * operation. A quiesced device cannot use the following functions:
+ * vfu_addr_to_sg, vfu_map_sg, vfu_unmap_sg, vfu_dma_read, and vfu_dma_write.
+ *
+ * The callback can return two values:
+ * 1) 0: this indicates that the device was quiesced. vfu_run_ctx then continues
+ *      to execute and when vfu_run_ctx returns to the caller the device is
+ *      unquiesced.
+ * 2) -1 with errno set to EBUSY: this indicates that the device cannot
+ *      immediately quiesce. In this case, vfu_run_ctx returns -1 with errno
+ *      set to EBUSY and future calls to vfu_run_ctx return the same. Until the
+ *      device quiesces it can continue operate as normal. The device indicates
+ *      that it quiesced by calling vfu_device_quiesced. When
+ *      vfu_device_quiesced returns the device is no longer quiesced.
+ *
+ * A quiesced device should expect for any of the following callbacks to be
+ * executed:
+ * vfu_dma_register_cb_t, vfu_unregister_cb_t, vfu_reset_cb_t, and transition.
+ * These callbacks are only called after the device has been quiesced.
+ *
+ * @vfu_ctx: the libvfio-user context
+ *
+ * @returns: 0 on success, -1 on failure with errno set.
+ */
+typedef int (vfu_device_quiesce_cb_t)(vfu_ctx_t *vfu_ctx);
+
+/**
+ * Sets up the device quiesce callback.
+ *
+ * @vfu_ctx: the libvfio-user context
+ * @quiesce_cb: device quiesce callback
+ */
+void
+vfu_setup_device_quiesce_cb(vfu_ctx_t *vfu_ctx,
+                            vfu_device_quiesce_cb_t *quiesce_cb);
+
+/*
+ * Called by the device to complete a pending quiesce operation. After the
+ * function returns the device is unquiesced.
+ *
+ * @vfu_ctx: the libvfio-user context
+ * @quiesce_errno: 0 for success or errno in case the device fails to quiesce,
+ *                 in which case the operation requiring the quiesce is failed
+ *                 and the device is reset.
+ *
+ * @returns 0 on success, or -1 on failure. Sets errno.
+ */
+int
+vfu_device_quiesced(vfu_ctx_t *vfu_ctx, int quiesce_errno);
+
+/*
  * Callback function that is called when the device must be reset.
  */
 typedef int (vfu_reset_cb_t)(vfu_ctx_t *vfu_ctx, vfu_reset_type_t type);
@@ -414,19 +470,15 @@ typedef struct vfu_dma_info {
 typedef void (vfu_dma_register_cb_t)(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info);
 
 /*
- * Function that is called when the guest unregisters a DMA region. The device
- * must release all references to that region before the callback returns.
- * This is required if you want to be able to access guest memory directly via
- * a mapping.
- *
- * The callback should return 0 on success, -1 with errno set on failure
- * (although unregister should not fail: this will not stop a guest from
- * unregistering the region).
+ * Function that is called when the guest unregisters a DMA region.  This
+ * callback is required if you want to be able to access guest memory directly
+ * via a mapping. The device must release all references to that region before
+ * the callback returns.
  *
  * @vfu_ctx: the libvfio-user context
  * @info: the DMA info
  */
-typedef int (vfu_dma_unregister_cb_t)(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info);
+typedef void (vfu_dma_unregister_cb_t)(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info);
 
 /**
  * Set up device DMA registration callbacks. When libvfio-user is notified of a
@@ -502,16 +554,6 @@ typedef struct {
      *
      * The callback should return -1 on error, setting errno.
      *
-     * When operating in non-blocking mode (LIBVFIO_USER_FLAG_ATTACH_NB was
-     * passed to vfu_create_ctx) and -1 is returned with errno set to EBUSY,
-     * transitioning to the new state becomes asynchronous: libvfio-user does
-     * not send a response to the client and does not process any new messages.
-     * Transitioning to the new device state is completed by calling
-     * vfu_migr_done. This behavior can be beneficial for devices whose
-     * threading model does not allow blocking.
-     *
-     * The user must not call functions vfu_dma_read or vfu_dma_write, doing so
-     * results in undefined behavior.
      *
      * TODO rename to vfu_migration_state_transition_callback
      * FIXME maybe we should create a single callback and pass the state?
@@ -579,18 +621,6 @@ typedef struct {
 
 } vfu_migration_callbacks_t;
 
-/*
- * Completes a pending migration state transition. Calling this function when
- * there is no pending migration state transition results in undefined
- * behavior.
- *
- * @vfu_ctx: the libvfio-user context
- * @reply_errno: 0 for success or errno on error.
- */
-void
-vfu_migr_done(vfu_ctx_t *vfu_ctx, int reply_errno);
-
-
 #ifndef VFIO_DEVICE_STATE_STOP
 
 #define VFIO_DEVICE_STATE_STOP (0)
diff --git a/lib/dma.c b/lib/dma.c
index eed4370..cf63177 100644
--- a/lib/dma.c
+++ b/lib/dma.c
@@ -156,7 +156,6 @@ MOCK_DEFINE(dma_controller_remove_region)(dma_controller_t *dma,
 {
     int idx;
     dma_memory_region_t *region;
-    int err;
 
     assert(dma != NULL);
 
@@ -167,13 +166,8 @@ MOCK_DEFINE(dma_controller_remove_region)(dma_controller_t *dma,
             continue;
         }
 
-        err = dma_unregister == NULL ? 0 : dma_unregister(data, &region->info);
-        if (err != 0) {
-            err = errno;
-            vfu_log(dma->vfu_ctx, LOG_ERR,
-                   "failed to dma_unregister() DMA region [%p, %p): %m",
-                   region->info.iova.iov_base, iov_end(&region->info.iova));
-            return ERROR_INT(err);
+        if (dma_unregister != NULL) {
+            dma_unregister(data, &region->info);
         }
 
         assert(region->refcnt == 0);
@@ -201,7 +195,6 @@ dma_controller_remove_all_regions(dma_controller_t *dma,
 
     for (i = 0; i < dma->nregions; i++) {
         dma_memory_region_t *region = &dma->regions[i];
-        int err;
 
         vfu_log(dma->vfu_ctx, LOG_DEBUG, "removing DMA region "
                 "iova=[%p, %p) vaddr=%p mapping=[%p, %p)",
@@ -209,12 +202,8 @@ dma_controller_remove_all_regions(dma_controller_t *dma,
                 region->info.vaddr,
                 region->info.mapping.iov_base, iov_end(&region->info.mapping));
 
-        err = dma_unregister == NULL ? 0 : dma_unregister(data, &region->info);
-        if (err != 0) {
-            err = errno;
-            vfu_log(dma->vfu_ctx, LOG_ERR,
-                   "failed to dma_unregister() DMA region [%p, %p): %m",
-                   region->info.iova.iov_base, iov_end(&region->info.iova));
+        if (dma_unregister != NULL) {
+            dma_unregister(data, &region->info);
         }
 
         if (region->info.vaddr != NULL) {
diff --git a/lib/libvfio-user.c b/lib/libvfio-user.c
index b2ffac6..44641c2 100644
--- a/lib/libvfio-user.c
+++ b/lib/libvfio-user.c
@@ -58,7 +58,8 @@
 #include "private.h"
 #include "tran_sock.h"
 
-static void vfu_reset_ctx(vfu_ctx_t *vfu_ctx, const char *reason);
+static int
+vfu_reset_ctx(vfu_ctx_t *vfu_ctx, int reason);
 
 EXPORT void
 vfu_log(vfu_ctx_t *vfu_ctx, int level, const char *fmt, ...)
@@ -195,6 +196,14 @@ dump_buffer(const char *prefix UNUSED, const char *buf UNUSED,
 #endif
 }
 
+static bool
+access_needs_quiesce(const vfu_ctx_t *vfu_ctx, size_t region_index,
+                     uint64_t offset)
+{
+    return access_migration_needs_quiesce(vfu_ctx, region_index, offset)
+           || access_is_pci_cap_exp(vfu_ctx, region_index, offset);
+}
+
 static ssize_t
 region_access(vfu_ctx_t *vfu_ctx, size_t region_index, char *buf,
               size_t count, uint64_t offset, bool is_write)
@@ -335,18 +344,6 @@ handle_region_access(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 
     ret = region_access(vfu_ctx, in_ra->region, buf, in_ra->count,
                         in_ra->offset, msg->hdr.cmd == VFIO_USER_REGION_WRITE);
-    if (ret == -1 && in_ra->region == VFU_PCI_DEV_MIGR_REGION_IDX
-        && errno == EBUSY && (vfu_ctx->flags & LIBVFIO_USER_FLAG_ATTACH_NB)) {
-        /*
-         * We don't support async behavior for the non-blocking mode simply
-         * because we don't have a use case yet, the only user of migration
-         * is SPDK and it operates in non-blocking mode. We don't know the
-         * implications of enabling this in blocking mode as we haven't looked
-         * at the details.
-         */
-        vfu_ctx->migr_trans_pending = true;
-        return 0;
-    }
     if (ret != in_ra->count) {
         vfu_log(vfu_ctx, LOG_ERR, "failed to %s %#lx-%#lx: %m",
                 msg->hdr.cmd == VFIO_USER_REGION_WRITE ? "write" : "read",
@@ -1086,13 +1083,16 @@ get_request_header(vfu_ctx_t *vfu_ctx, vfu_msg_t **msgp)
             return -1;
 
         case ENOMSG:
-            vfu_reset_ctx(vfu_ctx, "closed");
-            return ERROR_INT(ENOTCONN);
-
         case ECONNRESET:
-            vfu_reset_ctx(vfu_ctx, "reset");
+            vfu_log(vfu_ctx, LOG_DEBUG, "failed to receive request header: %m");
+            ret = vfu_reset_ctx(vfu_ctx, errno);
+            if (ret < 0) {
+                if (errno != EBUSY) {
+                    vfu_log(vfu_ctx, LOG_WARNING, "failed to reset context: %m");
+                }
+                return ret;
+            }
             return ERROR_INT(ENOTCONN);
-
         default:
             vfu_log(vfu_ctx, LOG_ERR, "failed to receive request: %m");
             return -1;
@@ -1179,6 +1179,42 @@ MOCK_DEFINE(should_exec_command)(vfu_ctx_t *vfu_ctx, uint16_t cmd)
     return true;
 }
 
+static bool
+command_needs_quiesce(vfu_ctx_t *vfu_ctx, const vfu_msg_t *msg)
+{
+    struct vfio_user_region_access *reg;
+
+    if (vfu_ctx->quiesce == NULL) {
+        return false;
+    }
+
+    switch (msg->hdr.cmd) {
+    case VFIO_USER_DMA_MAP:
+    case VFIO_USER_DMA_UNMAP:
+        return vfu_ctx->dma != NULL;
+
+    case VFIO_USER_DEVICE_RESET:
+        return true;
+
+    case VFIO_USER_REGION_WRITE:
+        if (msg->in_size < sizeof(*reg)) {
+            /*
+             * bad request, it will be eventually failed by
+             * handle_region_access
+             *
+             */
+            return false;
+        }
+        reg = msg->in_data;
+        if (access_needs_quiesce(vfu_ctx, reg->region, reg->offset)) {
+            return true;
+        }
+        break;
+    }
+
+    return false;
+}
+
 int
 exec_command(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
@@ -1259,11 +1295,14 @@ do_reply(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg, int reply_errno)
     if (ret < 0) {
         vfu_log(vfu_ctx, LOG_ERR, "failed to reply: %m");
 
-        if (errno == ECONNRESET) {
-            vfu_reset_ctx(vfu_ctx, "reset");
-            errno = ENOTCONN;
-        } else if (errno == ENOMSG) {
-            vfu_reset_ctx(vfu_ctx, "closed");
+        if (errno == ECONNRESET || errno == ENOMSG) {
+            ret = vfu_reset_ctx(vfu_ctx, errno);
+            if (ret < 0) {
+                if (errno != EBUSY) {
+                    vfu_log(vfu_ctx, LOG_WARNING, "failed to reset context: %m");
+                }
+                return ret;
+            }
             errno = ENOTCONN;
         }
     }
@@ -1278,53 +1317,68 @@ do_reply(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg, int reply_errno)
  * possibly reply.
  */
 static int
-process_request(vfu_ctx_t *vfu_ctx)
+process_request(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
-    vfu_msg_t *msg = NULL;
     int ret;
 
     assert(vfu_ctx != NULL);
 
-    ret = get_request_header(vfu_ctx, &msg);
+    if (msg == NULL) {
+        ret = get_request_header(vfu_ctx, &msg);
 
-    if (ret < 0) {
-        return ret;
-    }
+        if (ret < 0) {
+            return ret;
+        }
 
-    if (!is_valid_header(vfu_ctx, msg)) {
-        ret = ERROR_INT(EINVAL);
-        goto out;
-    }
+        if (!is_valid_header(vfu_ctx, msg)) {
+            ret = ERROR_INT(EINVAL);
+            goto out;
+        }
 
-    msg->in_size = msg->hdr.msg_size - sizeof(msg->hdr);
+        msg->in_size = msg->hdr.msg_size - sizeof(msg->hdr);
 
-    if (msg->in_size > 0) {
-        ret = vfu_ctx->tran->recv_body(vfu_ctx, msg);
+        if (msg->in_size > 0) {
+            ret = vfu_ctx->tran->recv_body(vfu_ctx, msg);
 
-        if (ret < 0) {
+            if (ret < 0) {
+                goto out;
+            }
+        }
+
+        if (!should_exec_command(vfu_ctx, msg->hdr.cmd)) {
+            ret = ERROR_INT(EINVAL);
             goto out;
         }
-    }
 
-    if (!should_exec_command(vfu_ctx, msg->hdr.cmd)) {
-        ret = ERROR_INT(EINVAL);
-        goto out;
+        if (command_needs_quiesce(vfu_ctx, msg)) {
+            vfu_log(vfu_ctx, LOG_DEBUG, "quiescing device");
+            ret = vfu_ctx->quiesce(vfu_ctx);
+            if (ret < 0) {
+                if (errno == EBUSY) {
+                    vfu_log(vfu_ctx, LOG_DEBUG, "device will quiesce asynchronously");
+                    vfu_ctx->pending.state = VFU_CTX_PENDING_MSG;
+                    vfu_ctx->pending.msg = msg;
+                    /* NB the message is freed in vfu_device_quiesced */
+                    return ret;
+                } else {
+                    vfu_log(vfu_ctx, LOG_DEBUG, "device failed to quiesce: %m");
+                }
+                goto out;
+            } else {
+                vfu_log(vfu_ctx, LOG_DEBUG, "device quiesced immediately");
+            }
+        }
     }
+    errno = 0;
 
     ret = exec_command(vfu_ctx, msg);
 
-    if (ret < 0) {
+    if (ret < 0 && errno != EBUSY) {
         vfu_log(vfu_ctx, LOG_ERR, "msg%#hx: cmd %d failed: %m", msg->hdr.msg_id,
                 msg->hdr.cmd);
     }
 
 out:
-    if (vfu_ctx->migr_trans_pending) {
-        assert(ret == 0);
-        vfu_ctx->migr_trans_msg = msg;
-        /* NB the message is freed in vfu_migr_done */
-        return 0;
-    }
     if (msg->hdr.flags.no_reply) {
         /*
          * A failed client request is not a failure of process_request() itself.
@@ -1333,7 +1387,6 @@ out:
     } else {
         ret = do_reply(vfu_ctx, msg, ret == 0 ? 0 : errno);
     }
-
     free_msg(vfu_ctx, msg);
     return ret;
 }
@@ -1426,16 +1479,17 @@ vfu_run_ctx(vfu_ctx_t *vfu_ctx)
     assert(vfu_ctx != NULL);
 
     if (!vfu_ctx->realized) {
+        vfu_log(vfu_ctx, LOG_DEBUG, "device not realized");
         return ERROR_INT(EINVAL);
     }
 
     blocking = !(vfu_ctx->flags & LIBVFIO_USER_FLAG_ATTACH_NB);
 
     do {
-        if (vfu_ctx->migr_trans_pending) {
+        if (vfu_ctx->pending.state != VFU_CTX_PENDING_NONE) {
             return ERROR_INT(EBUSY);
         }
-        err = process_request(vfu_ctx);
+        err = process_request(vfu_ctx, NULL);
 
         if (err == 0) {
             reqs_processed++;
@@ -1462,15 +1516,14 @@ free_sparse_mmap_areas(vfu_ctx_t *vfu_ctx)
 }
 
 static void
-vfu_reset_ctx(vfu_ctx_t *vfu_ctx, const char *reason)
+vfu_reset_ctx_quiesced(vfu_ctx_t *vfu_ctx)
 {
-    vfu_log(vfu_ctx, LOG_INFO, "%s: %s", __func__,  reason);
-
     if (vfu_ctx->dma != NULL) {
         dma_controller_remove_all_regions(vfu_ctx->dma, vfu_ctx->dma_unregister,
                                           vfu_ctx);
     }
 
+    /* FIXME what happens if the device reset callback fails? */
     do_device_reset(vfu_ctx, VFU_RESET_LOST_CONN);
 
     if (vfu_ctx->irqs != NULL) {
@@ -1482,15 +1535,38 @@ vfu_reset_ctx(vfu_ctx_t *vfu_ctx, const char *reason)
     }
 }
 
+static int
+vfu_reset_ctx(vfu_ctx_t *vfu_ctx, int reason)
+{
+    vfu_log(vfu_ctx, LOG_INFO, "%s: %s", __func__,  strerror(reason));
+
+    if (vfu_ctx->quiesce != NULL
+        && vfu_ctx->pending.state == VFU_CTX_PENDING_NONE) {
+        int ret = vfu_ctx->quiesce(vfu_ctx);
+        if (ret < 0) {
+            if (errno == EBUSY) {
+                vfu_ctx->pending.state = VFU_CTX_PENDING_CTX_RESET;
+                return ret;
+            }
+            vfu_log(vfu_ctx, LOG_ERR, "failed to quiesce device: %m");
+            return ret;
+        }
+    }
+    vfu_reset_ctx_quiesced(vfu_ctx);
+    return 0;
+}
+
 EXPORT void
 vfu_destroy_ctx(vfu_ctx_t *vfu_ctx)
 {
-
     if (vfu_ctx == NULL) {
         return;
     }
 
-    vfu_reset_ctx(vfu_ctx, "destroyed");
+    vfu_ctx->quiesce = NULL;
+    if (vfu_reset_ctx(vfu_ctx, ESHUTDOWN) < 0) {
+        vfu_log(vfu_ctx, LOG_WARNING, "failed to reset context: %m");
+    }
 
     free(vfu_ctx->uuid);
     free(vfu_ctx->pci.config_space);
@@ -1548,6 +1624,7 @@ vfu_create_ctx(vfu_trans_t trans, const char *path, int flags, void *pvt,
     vfu_ctx->pvt = pvt;
     vfu_ctx->flags = flags;
     vfu_ctx->log_level = LOG_ERR;
+    vfu_ctx->pci_cap_exp_off = -1;
 
     vfu_ctx->uuid = strdup(path);
     if (vfu_ctx->uuid == NULL) {
@@ -1784,6 +1861,13 @@ vfu_setup_device_reset_cb(vfu_ctx_t *vfu_ctx, vfu_reset_cb_t *reset)
     return 0;
 }
 
+EXPORT void
+vfu_setup_device_quiesce_cb(vfu_ctx_t *vfu_ctx, vfu_device_quiesce_cb_t *quiesce)
+{
+    assert(vfu_ctx != NULL);
+    vfu_ctx->quiesce = quiesce;
+}
+
 EXPORT int
 vfu_setup_device_dma(vfu_ctx_t *vfu_ctx, vfu_dma_register_cb_t *dma_register,
                      vfu_dma_unregister_cb_t *dma_unregister)
@@ -1874,6 +1958,8 @@ vfu_map_sg(vfu_ctx_t *vfu_ctx, dma_sg_t *sg, struct iovec *iov, int cnt,
         return ERROR_INT(EINVAL);
     }
 
+    assert(vfu_ctx->pending.state != VFU_CTX_PENDING_MSG);
+
     ret = dma_map_sg(vfu_ctx->dma, sg, iov, cnt);
     if (ret < 0) {
         return -1;
@@ -1951,11 +2037,10 @@ vfu_dma_transfer(vfu_ctx_t *vfu_ctx, enum vfio_user_command cmd,
 
         if (ret < 0) {
             ret = errno;
-            if (ret == ENOMSG) {
-                vfu_reset_ctx(vfu_ctx, "closed");
-                ret = ENOTCONN;
-            } else if (errno == ECONNRESET) {
-                vfu_reset_ctx(vfu_ctx, "reset");
+            if (ret == ENOMSG || ret == ECONNRESET) {
+                if (vfu_reset_ctx(vfu_ctx, ret) < 0) {
+                    vfu_log(vfu_ctx, LOG_WARNING, "failed to reset context: %m");
+                }
                 ret = ENOTCONN;
             }
             free(rbuf);
@@ -1987,14 +2072,14 @@ vfu_dma_transfer(vfu_ctx_t *vfu_ctx, enum vfio_user_command cmd,
 EXPORT int
 vfu_dma_read(vfu_ctx_t *vfu_ctx, dma_sg_t *sg, void *data)
 {
-    assert(!vfu_ctx->migr_trans_pending);
+    assert(vfu_ctx->pending.state == VFU_CTX_PENDING_NONE);
     return vfu_dma_transfer(vfu_ctx, VFIO_USER_DMA_READ, sg, data);
 }
 
 EXPORT int
 vfu_dma_write(vfu_ctx_t *vfu_ctx, dma_sg_t *sg, void *data)
 {
-    assert(!vfu_ctx->migr_trans_pending);
+    assert(vfu_ctx->pending.state == VFU_CTX_PENDING_NONE);
     return vfu_dma_transfer(vfu_ctx, VFIO_USER_DMA_WRITE, sg, data);
 }
 
@@ -2004,19 +2089,44 @@ vfu_sg_is_mappable(vfu_ctx_t *vfu_ctx, dma_sg_t *sg)
     return dma_sg_is_mappable(vfu_ctx->dma, sg);
 }
 
-EXPORT void
-vfu_migr_done(vfu_ctx_t *vfu_ctx, int reply_errno)
+EXPORT int
+vfu_device_quiesced(vfu_ctx_t *vfu_ctx, int quiesce_errno)
 {
+    int ret;
+
     assert(vfu_ctx != NULL);
-    assert(vfu_ctx->migr_trans_pending);
 
-    if (!vfu_ctx->migr_trans_msg->hdr.flags.no_reply) {
-        do_reply(vfu_ctx, vfu_ctx->migr_trans_msg, reply_errno);
+    if (vfu_ctx->quiesce == NULL
+        || vfu_ctx->pending.state == VFU_CTX_PENDING_NONE) {
+        vfu_log(vfu_ctx, LOG_DEBUG,
+                "invalid call to quiesce callback, state=%d",
+                vfu_ctx->pending.state);
+        return ERROR_INT(EINVAL);
     }
-    free_msg(vfu_ctx, vfu_ctx->migr_trans_msg);
-    vfu_ctx->migr_trans_msg = NULL;
 
-    vfu_ctx->migr_trans_pending = false;
+    vfu_log(vfu_ctx, LOG_DEBUG, "device quiesced with error=%d", quiesce_errno);
+
+    if (quiesce_errno == 0) {
+        switch (vfu_ctx->pending.state) {
+        case VFU_CTX_PENDING_MSG:
+            ret = process_request(vfu_ctx, vfu_ctx->pending.msg);
+            break;
+        case VFU_CTX_PENDING_CTX_RESET:
+            vfu_reset_ctx_quiesced(vfu_ctx);
+            ret = 0;
+            break;
+        default:
+            assert(false);
+        }
+    } else {
+        ret = 0;
+        free_msg(vfu_ctx, vfu_ctx->pending.msg);
+    }
+
+    vfu_ctx->pending.msg = NULL;
+    vfu_ctx->pending.state = VFU_CTX_PENDING_NONE;
+
+    return ret;
 }
 
 /* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
diff --git a/lib/migration.c b/lib/migration.c
index c8b97fa..855e9e7 100644
--- a/lib/migration.c
+++ b/lib/migration.c
@@ -154,8 +154,9 @@ MOCK_DEFINE(migr_trans_to_valid_state)(vfu_ctx_t *vfu_ctx, struct migration *mig
                                        uint32_t device_state, bool notify)
 {
     if (notify) {
-        int ret = state_trans_notify(vfu_ctx, migr->callbacks.transition,
-                                     device_state);
+        int ret;
+        ret = state_trans_notify(vfu_ctx, migr->callbacks.transition,
+                                 device_state);
         if (ret != 0) {
             return ret;
         }
@@ -423,11 +424,6 @@ MOCK_DEFINE(migration_region_access_registers)(vfu_ctx_t *vfu_ctx, char *buf,
                 "migration: transition from state %s to state %s",
                  migr_states[old_device_state].name,
                  migr_states[*device_state].name);
-        } else if (errno == EBUSY) {
-            vfu_log(vfu_ctx, LOG_DEBUG,
-                "migration: transition from state %s to state %s deferred",
-                 migr_states[old_device_state].name,
-                 migr_states[*device_state].name);
         } else {
             vfu_log(vfu_ctx, LOG_ERR,
                 "migration: failed to transition from state %s to state %s",
@@ -561,4 +557,18 @@ migration_set_pgsize(struct migration *migr, size_t pgsize)
     return 0;
 }
 
+bool
+access_migration_needs_quiesce(const vfu_ctx_t *vfu_ctx, size_t region_index,
+                              uint64_t offset)
+{
+    /*
+     * Writing to the migration state register with an unaligned access won't
+     * trigger this check but that's not a problem because
+     * migration_region_access_registers will fail the access.
+     */
+    return region_index == VFU_PCI_DEV_MIGR_REGION_IDX
+           && vfu_ctx->migration != NULL
+           && offset == offsetof(struct vfio_user_migration_info, device_state);
+}
+
 /* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
diff --git a/lib/migration.h b/lib/migration.h
index ccb98aa..26fd744 100644
--- a/lib/migration.h
+++ b/lib/migration.h
@@ -71,6 +71,10 @@ MOCK_DECLARE(bool, vfio_migr_state_transition_is_valid, uint32_t from,
 MOCK_DECLARE(ssize_t, handle_device_state, vfu_ctx_t *vfu_ctx,
              struct migration *migr, uint32_t device_state, bool notify);
 
+bool
+access_migration_needs_quiesce(const vfu_ctx_t *vfu_ctx, size_t region_index,
+                              uint64_t offset);
+
 #endif /* LIB_VFIO_USER_MIGRATION_H */
 
 /* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
diff --git a/lib/migration_priv.h b/lib/migration_priv.h
index 340d507..bf7c7b0 100644
--- a/lib/migration_priv.h
+++ b/lib/migration_priv.h
@@ -127,9 +127,6 @@ MOCK_DECLARE(int, state_trans_notify, vfu_ctx_t *vfu_ctx,
              int (*fn)(vfu_ctx_t *, vfu_migr_state_t),
              uint32_t vfio_device_state);
 
-MOCK_DECLARE(ssize_t, migr_trans_to_valid_state, vfu_ctx_t *vfu_ctx,
-             struct migration *migr, uint32_t device_state, bool notify);
-
 #endif
 
 /* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
diff --git a/lib/pci_caps.c b/lib/pci_caps.c
index ec0dbd8..cd91914 100644
--- a/lib/pci_caps.c
+++ b/lib/pci_caps.c
@@ -623,11 +623,13 @@ vfu_pci_add_capability(vfu_ctx_t *vfu_ctx, size_t pos, int flags, void *data)
 
     if (flags & ~(VFU_CAP_FLAG_EXTENDED | VFU_CAP_FLAG_CALLBACK |
         VFU_CAP_FLAG_READONLY)) {
+        vfu_log(vfu_ctx, LOG_DEBUG, "bad flags %#x", flags);
         return ERROR_INT(EINVAL);
     }
 
     if ((flags & VFU_CAP_FLAG_CALLBACK) &&
         vfu_ctx->reg_info[VFU_PCI_DEV_CFG_REGION_IDX].cb == NULL) {
+        vfu_log(vfu_ctx, LOG_DEBUG, "no callback");
         return ERROR_INT(EINVAL);
     }
 
@@ -641,6 +643,7 @@ vfu_pci_add_capability(vfu_ctx_t *vfu_ctx, size_t pos, int flags, void *data)
         case VFU_PCI_TYPE_EXPRESS:
             break;
         default:
+            vfu_log(vfu_ctx, LOG_DEBUG, "bad PCI type %#x", vfu_ctx->pci.type);
             return ERROR_INT(EINVAL);
         }
 
@@ -669,6 +672,7 @@ vfu_pci_add_capability(vfu_ctx_t *vfu_ctx, size_t pos, int flags, void *data)
         cap.size = cap_size(vfu_ctx, data, extended);
 
         if (cap.off + cap.size >= pci_config_space_size(vfu_ctx)) {
+            vfu_log(vfu_ctx, LOG_DEBUG, "bad PCIe capability offset");
             return ERROR_INT(EINVAL);
         }
 
@@ -708,6 +712,9 @@ vfu_pci_add_capability(vfu_ctx_t *vfu_ctx, size_t pos, int flags, void *data)
         cap.size = cap_size(vfu_ctx, data, extended);
 
         if (cap.off + cap.size >= pci_config_space_size(vfu_ctx)) {
+                vfu_log(vfu_ctx, LOG_DEBUG,
+                        "PCI capability past end of config space, %#lx >= %#lx",
+                        cap.off + cap.size, pci_config_space_size(vfu_ctx));
             return ERROR_INT(EINVAL);
         }
 
@@ -730,6 +737,10 @@ vfu_pci_add_capability(vfu_ctx_t *vfu_ctx, size_t pos, int flags, void *data)
         vfu_ctx->pci.nr_caps++;
     }
 
+
+    if (cap.id == PCI_CAP_ID_EXP) {
+        vfu_ctx->pci_cap_exp_off = cap.off;
+    }
     return cap.off;
 }
 
@@ -832,4 +843,12 @@ vfu_pci_find_capability(vfu_ctx_t *vfu_ctx, bool extended, int cap_id)
     return vfu_pci_find_next_capability(vfu_ctx, extended, 0, cap_id);
 }
 
+bool
+access_is_pci_cap_exp(const vfu_ctx_t *vfu_ctx, size_t region_index,
+                      uint64_t offset)
+{
+    size_t _offset = vfu_ctx->pci_cap_exp_off + offsetof(struct pxcap, pxdc);
+    return region_index == VFU_PCI_DEV_CFG_REGION_IDX && offset == _offset;
+}
+
 /* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
diff --git a/lib/pci_caps.h b/lib/pci_caps.h
index 4bef85d..f4b0ff2 100644
--- a/lib/pci_caps.h
+++ b/lib/pci_caps.h
@@ -72,6 +72,10 @@ ssize_t
 pci_cap_access(vfu_ctx_t *ctx, char *buf, size_t count, loff_t offset,
                bool is_write);
 
+bool
+access_is_pci_cap_exp(const vfu_ctx_t *vfu_ctx, size_t region_index,
+                      uint64_t offset);
+
 #endif /* LIB_VFIO_USER_PCI_CAPS_H */
 
 /* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
diff --git a/lib/private.h b/lib/private.h
index 05d2fa4..2e9a769 100644
--- a/lib/private.h
+++ b/lib/private.h
@@ -149,11 +149,25 @@ struct pci_dev {
 
 struct dma_controller;
 
+enum vfu_ctx_pending_state {
+    VFU_CTX_PENDING_NONE,
+    VFU_CTX_PENDING_MSG,
+    VFU_CTX_PENDING_DEVICE_RESET,
+    VFU_CTX_PENDING_CTX_RESET
+};
+
+struct vfu_ctx_pending_info {
+    enum vfu_ctx_pending_state  state;
+    vfu_msg_t                   *msg;
+
+    /* when pending == VFU_CTX_PENDING_XXX_RESET */
+    uint32_t                migr_dev_state;
+};
+
 struct vfu_ctx {
     void                    *pvt;
     struct dma_controller   *dma;
-    vfu_reset_cb_t          *reset;
-    int           log_level;
+    int                     log_level;
     vfu_log_fn_t            *log;
     size_t                  nr_regions;
     vfu_reg_info_t          *reg_info;
@@ -162,20 +176,26 @@ struct vfu_ctx {
     void                    *tran_data;
     uint64_t                flags;
     char                    *uuid;
+
+    /* device callbacks */
+    vfu_device_quiesce_cb_t *quiesce;
+    vfu_reset_cb_t          *reset;
     vfu_dma_register_cb_t   *dma_register;
     vfu_dma_unregister_cb_t *dma_unregister;
 
     int                     client_max_fds;
     size_t                  client_max_data_xfer_size;
 
+    struct vfu_ctx_pending_info pending;
+
     struct migration        *migration;
-    bool                    migr_trans_pending;
-    vfu_msg_t               *migr_trans_msg;
 
     uint32_t                irq_count[VFU_DEV_NUM_IRQS];
     vfu_irqs_t              *irqs;
     bool                    realized;
     vfu_dev_type_t          dev_type;
+
+    ssize_t                 pci_cap_exp_off;
 };
 
 typedef struct ioeventfd {
@@ -228,6 +248,9 @@ MOCK_DECLARE(bool, cmd_allowed_when_stopped_and_copying, uint16_t cmd);
 
 MOCK_DECLARE(bool, should_exec_command, vfu_ctx_t *vfu_ctx, uint16_t cmd);
 
+MOCK_DECLARE(ssize_t, migr_trans_to_valid_state, vfu_ctx_t *vfu_ctx,
+             struct migration *migr, uint32_t device_state, bool notify);
+
 #endif /* LIB_VFIO_USER_PRIVATE_H */
 
 /* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
diff --git a/samples/gpio-pci-idio-16.c b/samples/gpio-pci-idio-16.c
index d773067..4872d3b 100644
--- a/samples/gpio-pci-idio-16.c
+++ b/samples/gpio-pci-idio-16.c
@@ -130,10 +130,9 @@ dma_register(UNUSED vfu_ctx_t *vfu_ctx, UNUSED vfu_dma_info_t *info)
 {
 }
 
-static int
+static void
 dma_unregister(UNUSED vfu_ctx_t *vfu_ctx, UNUSED vfu_dma_info_t *info)
 {
-    return 0;
 }
 
 int
diff --git a/samples/server.c b/samples/server.c
index 416d5c8..a43dc3c 100644
--- a/samples/server.c
+++ b/samples/server.c
@@ -170,7 +170,7 @@ dma_register(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info)
     server_data->regions[idx].prot = info->prot;
 }
 
-static int
+static void
 dma_unregister(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info)
 {
     struct server_data *server_data = vfu_get_private(vfu_ctx);
@@ -181,11 +181,8 @@ dma_unregister(vfu_ctx_t *vfu_ctx, vfu_dma_info_t *info)
             server_data->regions[idx].iova.iov_base == info->iova.iov_base) {
             server_data->regions[idx].iova.iov_base = NULL;
             server_data->regions[idx].iova.iov_len = 0;
-            return 0;
         }
     }
-
-    return ERROR_INT(EINVAL);
 }
 
 static void
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,
-- 
cgit v1.1