From b0b8db12a5d8ea40330cbc47371f5f1e8f00bdb3 Mon Sep 17 00:00:00 2001
From: Pavel Minaev <pminaev@microsoft.com>
Date: Fri, 20 May 2022 15:12:24 -0700
Subject: [PATCH] Add test for wait_for_client.cancel()

---
 src/debugpy/public_api.py    |  4 ++--
 tests/debugpy/test_attach.py | 39 ++++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/src/debugpy/public_api.py b/src/debugpy/public_api.py
index 3e282065..36f0b513 100644
--- a/src/debugpy/public_api.py
+++ b/src/debugpy/public_api.py
@@ -28,7 +28,7 @@ def _api(cancelable=False):
             from debugpy.server import api
 
             wrapped = getattr(api, f.__name__)
-            wrapped(*args, **kwargs)
+            return wrapped(*args, **kwargs)
 
         if cancelable:
 
@@ -36,7 +36,7 @@ def _api(cancelable=False):
                 from debugpy.server import api
 
                 wrapped = getattr(api, f.__name__)
-                wrapped.cancel(*args, **kwargs)
+                return wrapped.cancel(*args, **kwargs)
 
             wrapper.cancel = cancel
 
diff --git a/tests/debugpy/test_attach.py b/tests/debugpy/test_attach.py
index 346fdfb4..a39524e2 100644
--- a/tests/debugpy/test_attach.py
+++ b/tests/debugpy/test_attach.py
@@ -209,3 +209,42 @@ def test_attach_pid_client(pyfile, target, pid_type):
         )
         session2.scratchpad["exit"] = True
         session2.request_continue()
+
+
+def test_cancel_wait(pyfile):
+    @pyfile
+    def code_to_debug():
+        import debuggee
+        import debugpy
+        import sys
+        import threading
+        import time
+
+        from debuggee import backchannel
+
+        def cancel():
+            time.sleep(1)
+            debugpy.wait_for_client.cancel()
+
+        _, host, port = sys.argv
+        port = int(port)
+        debugpy.listen(address=(host, port))
+        threading.Thread(target=cancel).start()
+        debugpy.wait_for_client()
+        backchannel.send("exit")
+
+    with debug.Session() as session:
+        host, port = runners.attach_connect.host, runners.attach_connect.port
+        session.config.update({"connect": {"host": host, "port": port}})
+        session.expected_exit_code = None
+
+        backchannel = session.open_backchannel()
+        session.spawn_debuggee(
+            [
+                code_to_debug,
+                host,
+                port,
+            ]
+        )
+
+        assert backchannel.receive() == "exit"