Allow transport tests' server to specify which class to use

1 files changed, 7 insertions, 2 deletions

diff --git a/tests/test_transport.py b/tests/test_transport.py
index 7a361080..d7704af6 100644
--- a/tests/test_transport.py
+++ b/tests/test_transport.py
@@ -1199,6 +1199,7 @@ def server(
     connect=None,
     pubkeys=None,
     catch_error=False,
+    transport_factory=None,
 ):
     """
     SSH server contextmanager for testing.
@@ -1219,6 +1220,8 @@ def server(
     :param catch_error:
         Whether to capture connection errors & yield from contextmanager.
         Necessary for connection_time exception testing.
+    :param transport_factory:
+        Like the same-named param in SSHClient: which Transport class to use.
     """
     if init is None:
         init = {}
@@ -1231,8 +1234,10 @@ def server(
     socks = LoopSocket()
     sockc = LoopSocket()
     sockc.link(socks)
-    tc = Transport(sockc, **dict(init, **client_init))
-    ts = Transport(socks, **dict(init, **server_init))
+    if transport_factory is None:
+        transport_factory = Transport
+    tc = transport_factory(sockc, **dict(init, **client_init))
+    ts = transport_factory(socks, **dict(init, **server_init))
 
     if hostkey is None:
         hostkey = RSAKey.from_private_key_file(_support("rsa.key"))