From b18b9ced293f580d588f321d41cda9dd322d417b Mon Sep 17 00:00:00 2001
From: Eric Snow <ericsnowcurrently@gmail.com>
Date: Mon, 2 Apr 2018 14:37:30 -0600
Subject: [PATCH] Wait for commands instead of just expecting them.

---
 tests/highlevel/__init__.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tests/highlevel/__init__.py b/tests/highlevel/__init__.py
index 86da125b..1587cbc4 100644
--- a/tests/highlevel/__init__.py
+++ b/tests/highlevel/__init__.py
@@ -287,8 +287,8 @@ class PyDevdLifecycle(object):
 
     @contextlib.contextmanager
     def _wait_for_initialized(self):
-        with self._fix.expect_command(CMD_REDIRECT_OUTPUT):
-            with self._fix.expect_command(CMD_RUN):
+        with self._fix.wait_for_command(CMD_REDIRECT_OUTPUT):
+            with self._fix.wait_for_command(CMD_RUN):
                 yield
 
     def _initialize(self):
@@ -529,6 +529,13 @@ class PyDevdFixture(FixtureBase):
         if self._hidden:
             self.msgs.next_request()
 
+    @contextlib.contextmanager
+    def wait_for_command(self, cmdid, *args, **kwargs):
+        with self.fake.wait_for_command(cmdid, *args, **kwargs):
+            yield
+        if self._hidden:
+            self.msgs.next_request()
+
     def set_response(self, cmdid, payload, **kwargs):
         self.fake.add_pending_response(cmdid, payload, **kwargs)
         if self._hidden: