Tests refactor

tests/system/test_post_mortem.py

@@ -1,6 +1,6 @@
 import madbg
 
-from .utils import run_in_process, run_script_in_process, JOIN_TIMEOUT, run_client
+from .utils import run_in_process, run_script_in_process, run_client
 
 
 def post_mortem_script(port):
@@ -11,8 +11,5 @@ def post_mortem_script(port):
 
 
 def test_post_mortem(port, start_debugger_with_ctty):
-    debugger_future = run_script_in_process(post_mortem_script, start_debugger_with_ctty, port)
-    assert not debugger_future.done()
-    client_future = run_in_process(run_client, port, b'c\n')
-    debugger_future.result(JOIN_TIMEOUT)
-    client_future.result(JOIN_TIMEOUT)
+    with run_script_in_process(post_mortem_script, start_debugger_with_ctty, port):
+        run_in_process(run_client, port, b'c\n').finish()

tests/system/test_run_with_debugging.py

@@ -1,8 +1,7 @@
-import time
 from pytest import raises
 from madbg import run_with_debugging
 
-from .utils import run_script_in_process, JOIN_TIMEOUT, SCRIPTS_PATH, run_in_process, run_client
+from .utils import run_script_in_process, SCRIPTS_PATH, run_in_process, run_client
 
 
 def run_divide_with_zero_with_debugging_script(port, post_mortem, set_trace):
@@ -11,20 +10,13 @@ def run_divide_with_zero_with_debugging_script(port, post_mortem, set_trace):
 
 
 def test_run_with_debugging_with_post_mortem(port, start_debugger_with_ctty):
-    debugger_future = run_script_in_process(run_divide_with_zero_with_debugging_script, start_debugger_with_ctty, port,
-                                            set_trace=False, post_mortem=True)
-    time.sleep(1)
-    assert not debugger_future.done()
-    client_future = run_in_process(run_client, port, b'c\n')
     with raises(ZeroDivisionError):
-        debugger_future.result(JOIN_TIMEOUT)
-    client_future.result(JOIN_TIMEOUT)
+        with run_script_in_process(run_divide_with_zero_with_debugging_script, start_debugger_with_ctty, port,
+                                   set_trace=False, post_mortem=True):
+            run_in_process(run_client, port, b'c\n').finish()
 
 
 def test_run_with_debugging_with_set_trace(port, start_debugger_with_ctty):
-    debugger_future = run_script_in_process(run_divide_with_zero_with_debugging_script, start_debugger_with_ctty, port,
-                                            set_trace=True, post_mortem=False)
-    assert not debugger_future.done()
-    client_future = run_in_process(run_client, port, b'n\nn\nyo = 0\nc\n')
-    debugger_future.result(JOIN_TIMEOUT)
-    client_future.result(JOIN_TIMEOUT)
+    with run_script_in_process(run_divide_with_zero_with_debugging_script, start_debugger_with_ctty, port,
+                                            set_trace=True, post_mortem=False):
+        run_in_process(run_client, port, b'n\nn\nyo = 0\nc\n').finish()

tests/system/test_set_trace.py

@@ -1,13 +1,14 @@
 import madbg
-from unittest.mock import Mock
 from pytest import raises
 
 from madbg.debugger import RemoteIPythonDebugger
 
-from .utils import run_in_process, run_script_in_process, JOIN_TIMEOUT, run_client
+from .utils import run_in_process, run_script_in_process, run_client
 
 
-def set_trace_script(port, times=1):
+def set_trace_script(port, times=1, debugger_fails=False):
+    if debugger_fails:
+        RemoteIPythonDebugger.__init__ = lambda *a, **k: 1 / 0
     for _ in range(times):
         madbg.set_trace(port=port)
 
@@ -22,38 +23,28 @@ def set_trace_and_expect_var_to_change_script(port) -> bool:
 
 
 def test_set_trace(port, start_debugger_with_ctty):
-    debugger_future = run_script_in_process(set_trace_and_expect_var_to_change_script, start_debugger_with_ctty, port)
-    client_future = run_in_process(run_client, port, b'value_to_change += 1\nc\n')
-    assert debugger_future.result(JOIN_TIMEOUT)
-    client_output = client_future.result(JOIN_TIMEOUT)
-    assert b'Closing connection' in client_output
+    with run_script_in_process(set_trace_and_expect_var_to_change_script, start_debugger_with_ctty, port):
+        assert b'Closing connection' in run_in_process(run_client, port, b'value_to_change += 1\nc\n').finish().get(0)
 
 
 def test_set_trace_and_connect_twice(port, start_debugger_with_ctty):
-    debugger_future = run_script_in_process(set_trace_script, start_debugger_with_ctty, port, 2)
-    assert b'Closing connection' in run_in_process(run_client, port, b'q\n').result(JOIN_TIMEOUT)
-    assert b'Closing connection' in run_in_process(run_client, port, b'q\n').result(JOIN_TIMEOUT)
-    debugger_future.result(JOIN_TIMEOUT)
+    with run_script_in_process(set_trace_script, start_debugger_with_ctty, port, 2):
+        assert b'Closing connection' in run_in_process(run_client, port, b'q\n').finish().get(0)
+        assert b'Closing connection' in run_in_process(run_client, port, b'q\n').finish().get(0)
 
 
 def test_set_trace_twice_and_continue(port, start_debugger_with_ctty):
-    debugger_future = run_script_in_process(set_trace_script, start_debugger_with_ctty, port, 2)
-    assert b'Closing connection' in run_in_process(run_client, port, b'c\nq\n').result(JOIN_TIMEOUT)
-    debugger_future.result(JOIN_TIMEOUT)
+    with run_script_in_process(set_trace_script, start_debugger_with_ctty, port, 2):
+        assert b'Closing connection' in run_in_process(run_client, port, b'c\nq\n').finish().get(0)
 
 
 def test_set_trace_and_quit_debugger(port, start_debugger_with_ctty):
-    debugger_future = run_script_in_process(set_trace_script, start_debugger_with_ctty, port)
-    client_future = run_in_process(run_client, port, b'q\n')
-    debugger_future.result(JOIN_TIMEOUT)
-    client_future.result(JOIN_TIMEOUT)
+    with run_script_in_process(set_trace_script, start_debugger_with_ctty, port):
+        run_in_process(run_client, port, b'q\n').finish()
 
 
 def test_set_trace_with_failing_debugger(port, start_debugger_with_ctty, monkeypatch):
-    monkeypatch.setattr(RemoteIPythonDebugger, '__init__', Mock(side_effect=lambda *a, **k: 1 / 0))
-    debugger_future = run_script_in_process(set_trace_script, start_debugger_with_ctty, port)
-    client_future = run_in_process(run_client, port, b'bla\n')
+    with run_script_in_process(set_trace_script, start_debugger_with_ctty, port, debugger_fails=True) as script_result:
+        assert ZeroDivisionError.__name__.encode() in run_in_process(run_client, port, b'bla\n').finish().get(0)
     with raises(ZeroDivisionError):
-        debugger_future.result(JOIN_TIMEOUT)
-    client_output = client_future.result(JOIN_TIMEOUT)
-    assert ZeroDivisionError.__name__.encode() in client_output
+        script_result.get(0)

tests/system/test_set_trace_on_connect.py

@@ -1,7 +1,7 @@
 import time
 import madbg
 
-from .utils import run_in_process, run_script_in_process, JOIN_TIMEOUT, run_client
+from .utils import run_in_process, run_script_in_process, run_client
 
 
 def set_trace_on_connect_script(port) -> bool:
@@ -16,16 +16,12 @@ def set_trace_on_connect_script(port) -> bool:
 
 
 def test_set_trace_on_connect(port, start_debugger_with_ctty):
-    debugger_future = run_script_in_process(set_trace_on_connect_script, start_debugger_with_ctty, port)
-    # let the loop run a little
-    time.sleep(0.5)
-    assert not debugger_future.done()
-    # Test we can connect twice
-    run_in_process(run_client, port, b'q\n').result(JOIN_TIMEOUT)
-    client_future = run_in_process(run_client, port, b'conti = False\nc\n')
-    assert debugger_future.result(JOIN_TIMEOUT)
-    client_future.result(JOIN_TIMEOUT)
+    with run_script_in_process(set_trace_on_connect_script, start_debugger_with_ctty, port) as script_result:
+        # Test we can connect twice
+        run_in_process(run_client, port, b'q\n').finish()
+        run_in_process(run_client, port, b'conti = False\nc\n').finish()
+    assert script_result.get(0)
 
 
 def test_set_trace_on_connect_can_exit(port, start_debugger_with_ctty):
-    run_script_in_process(madbg.set_trace_on_connect, start_debugger_with_ctty, port=port).result(JOIN_TIMEOUT)
+    run_script_in_process(madbg.set_trace_on_connect, start_debugger_with_ctty, port=port).finish()

tests/system/utils.py

@@ -1,34 +1,59 @@
-import atexit
 import os
 import pty
 import select
 import socket
-from concurrent.futures import ProcessPoolExecutor
-from contextlib import closing
+import multiprocessing as mp
+from contextlib import closing, _GeneratorContextManager
+from functools import wraps
 from pathlib import Path
 
 from madbg import client
 from madbg.consts import STDIN_FILENO, STDOUT_FILENO, STDERR_FILENO
 from madbg.tty_utils import PTY
 
-JOIN_TIMEOUT = 5
+JOIN_TIMEOUT = 10
 CONNECT_TIMEOUT = 5
 SCRIPTS_PATH = Path(__file__).parent / 'scripts'
 
+# forked subprocesses don't run exitfuncs
+mp_context = mp.get_context("spawn")
 
+
+class FinishableGeneratorContextManager(_GeneratorContextManager):
+    def finish(self):
+        with self as result:
+            return result
+
+
+def finishable_contextmanager(func):
+    @wraps(func)
+    def helper(*args, **kwds):
+        return FinishableGeneratorContextManager(func, args, kwds)
+    return helper
+
+
+@finishable_contextmanager
 def run_in_process(func, *args, **kwargs):
-    return ProcessPoolExecutor(1).submit(func, *args, **kwargs)
+    pool = mp_context.Pool(1)
+    apply_result = pool.apply_async(func, args, kwargs)
+    pool.close()
+    try:
+        yield apply_result
+    except:
+        pool.terminate()
+        raise
+    else:
+        # Wait for the result and raise an error if failed
+        apply_result.get(JOIN_TIMEOUT)
+        pool.join()
 
 
 def _run_script(script, start_with_ctty, args, kwargs):
     """
     Meant to be called inside a python subprocess, do NOT call directly.
     """
     enter_pty(start_with_ctty)
-    result = script(*args, **kwargs)
-    # Python-spawned subprocesses do not call exit funcs - https://stackoverflow.com/q/34506638/2907819
-    atexit._run_exitfuncs()
-    return result
+    return script(*args, **kwargs)
 
 
 def run_script_in_process(script, start_with_ctty, *args, **kwargs):

tox.ini

@@ -9,12 +9,9 @@ python =
     3.8 = py38
     3.7 = py37
 
-
 [testenv]
 package = wheel
-deps =
-    pytest
-    pytest-xdist
+deps = pytest
 commands = pytest tests {posargs}
 
 [testenv:sdist]