Introduce client object in python tests

test/py/libvfio_user.py

@@ -32,6 +32,7 @@
 #
 
 from types import SimpleNamespace
+import collections.abc
 import ctypes as c
 import array
 import errno
@@ -685,25 +686,60 @@ def connect_sock():
     return sock
 
 
-def connect_client(ctx):
-    sock = connect_sock()
-
-    json = b'{ "capabilities": { "max_msg_fds": 8 } }'
-    # struct vfio_user_version
-    payload = struct.pack("HH%dsc" % len(json), LIBVFIO_USER_MAJOR,
-                          LIBVFIO_USER_MINOR, json, b'\0')
-    hdr = vfio_user_header(VFIO_USER_VERSION, size=len(payload))
-    sock.send(hdr + payload)
-    vfu_attach_ctx(ctx, expect=0)
-    payload = get_reply(sock, expect=0)
-    return sock
-
-
-def disconnect_client(ctx, sock):
-    sock.close()
-
-    # notice client closed connection
-    vfu_run_ctx(ctx, errno.ENOTCONN)
+class Client:
+    """Models a VFIO-user client connected to the server under test."""
+
+    def __init__(self, sock=None):
+        self.sock = sock
+        self.client_cmd_socket = None
+
+    def connect(self, ctx, capabilities={}):
+        self.sock = connect_sock()
+
+        effective_caps = {
+            "capabilities": {
+                "max_data_xfer_size": VFIO_USER_DEFAULT_MAX_DATA_XFER_SIZE,
+                "max_msg_fds": 8,
+                "twin_socket": False,
+            },
+        }
+
+        def update(target, overrides):
+            for k, v in overrides.items():
+                if isinstance(v, collections.abc.Mapping):
+                    target[k] = target.get(k, {})
+                    update(target[k], v)
+                else:
+                    target[k] = v
+
+        update(effective_caps, capabilities)
+        caps_json = json.dumps(effective_caps)
+
+        # struct vfio_user_version
+        payload = struct.pack("HH%dsc" % len(caps_json), LIBVFIO_USER_MAJOR,
+                              LIBVFIO_USER_MINOR, caps_json.encode(), b'\0')
+        hdr = vfio_user_header(VFIO_USER_VERSION, size=len(payload))
+        self.sock.send(hdr + payload)
+        vfu_attach_ctx(ctx, expect=0)
+        fds, payload = get_reply_fds(self.sock, expect=0)
+        self.client_cmd_socket = socket.socket(fileno=fds[0]) if fds else None
+        return self.sock
+
+    def disconnect(self, ctx):
+        self.sock.close()
+        self.sock = None
+        if self.client_cmd_socket is not None:
+            self.client_cmd_socket.close()
+            self.client_cmd_socket = None
+
+        # notice client closed connection
+        vfu_run_ctx(ctx, errno.ENOTCONN)
+
+
+def connect_client(*args, **kwargs):
+    client = Client()
+    client.connect(*args, **kwargs)
+    return client
 
 
 def get_reply(sock, expect=0):

test/py/test_destroy.py

@@ -35,10 +35,10 @@
 
 
 def setup_function(function):
-    global ctx, sock
+    global ctx, client
     ctx = prepare_ctx_for_dma()
     assert ctx is not None
-    sock = connect_client(ctx)
+    client = connect_client(ctx)
 
 
 def teardown_function(function):

test/py/test_device_get_info.py

@@ -49,27 +49,27 @@ def test_device_get_info():
 
     # test short write
 
-    sock = connect_client(ctx)
+    client = connect_client(ctx)
 
     payload = struct.pack("II", 0, 0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_GET_INFO, payload,
+    msg(ctx, client.sock, VFIO_USER_DEVICE_GET_INFO, payload,
         expect=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,
+    msg(ctx, client.sock, VFIO_USER_DEVICE_GET_INFO, payload,
         expect=errno.EINVAL)
 
     # valid with larger argsz
 
     payload = vfio_user_device_info(argsz=32, flags=0,
                                     num_regions=0, num_irqs=0)
 
-    result = msg(ctx, sock, VFIO_USER_DEVICE_GET_INFO, payload)
+    result = msg(ctx, client.sock, VFIO_USER_DEVICE_GET_INFO, payload)
 
     (argsz, flags, num_regions, num_irqs) = struct.unpack("IIII", result)
 
@@ -78,7 +78,7 @@ def test_device_get_info():
     assert num_regions == VFU_PCI_DEV_NUM_REGIONS
     assert num_irqs == VFU_DEV_NUM_IRQS
 
-    disconnect_client(ctx, sock)
+    client.disconnect(ctx)
 
     vfu_destroy_ctx(ctx)
 

test/py/test_device_get_irq_info.py

@@ -31,13 +31,13 @@
 import errno
 
 ctx = None
-sock = None
+client = None
 
 argsz = len(vfio_irq_info())
 
 
 def test_device_get_irq_info_setup():
-    global ctx, sock
+    global ctx, client
 
     ctx = vfu_create_ctx(flags=LIBVFIO_USER_FLAG_ATTACH_NB)
     assert ctx is not None
@@ -55,27 +55,27 @@ def test_device_get_irq_info_setup():
     ret = vfu_realize_ctx(ctx)
     assert ret == 0
 
-    sock = connect_client(ctx)
+    client = connect_client(ctx)
 
 
 def test_device_get_irq_info_bad_in():
     payload = struct.pack("II", 0, 0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload,
+    msg(ctx, client.sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload,
         expect=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,
+    msg(ctx, client.sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload,
         expect=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,
+    msg(ctx, client.sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload,
         expect=errno.EINVAL)
 
 
@@ -86,12 +86,12 @@ def test_device_get_irq_info():
     payload = vfio_irq_info(argsz=argsz + 16, flags=0, index=VFU_DEV_REQ_IRQ,
                             count=0)
 
-    msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload)
+    msg(ctx, client.sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload)
 
     payload = vfio_irq_info(argsz=argsz, flags=0, index=VFU_DEV_REQ_IRQ,
                             count=0)
 
-    result = msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload)
+    result = msg(ctx, client.sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload)
 
     info, _ = vfio_irq_info.pop_from_buffer(result)
 
@@ -103,7 +103,7 @@ def test_device_get_irq_info():
     payload = vfio_irq_info(argsz=argsz, flags=0, index=VFU_DEV_ERR_IRQ,
                             count=0)
 
-    result = msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload)
+    result = msg(ctx, client.sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload)
 
     info, _ = vfio_irq_info.pop_from_buffer(result)
 
@@ -115,7 +115,7 @@ def test_device_get_irq_info():
     payload = vfio_irq_info(argsz=argsz, flags=0, index=VFU_DEV_MSIX_IRQ,
                             count=0)
 
-    result = msg(ctx, sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload)
+    result = msg(ctx, client.sock, VFIO_USER_DEVICE_GET_IRQ_INFO, payload)
 
     info, _ = vfio_irq_info.pop_from_buffer(result)
 
@@ -126,7 +126,7 @@ def test_device_get_irq_info():
 
 
 def test_device_get_irq_info_cleanup():
-    disconnect_client(ctx, sock)
+    client.disconnect(ctx)
 
     vfu_destroy_ctx(ctx)