From ff30675a256294d55dd39bd5b38163ecb3138d7b Mon Sep 17 00:00:00 2001
From: Romuald Juchnowicz-Bierbasz <rbierbasz@gog.com>
Date: Thu, 27 Jun 2019 21:18:16 +0200
Subject: [PATCH] Do not invoke tick before handshake

---
 src/galaxy/api/jsonrpc.py |  1 +
 src/galaxy/api/plugin.py  | 20 ++++++++++++--------
 tests/test_internal.py    | 13 ++++++++++++-
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/src/galaxy/api/jsonrpc.py b/src/galaxy/api/jsonrpc.py
index 37d0e17..87bff71 100644
--- a/src/galaxy/api/jsonrpc.py
+++ b/src/galaxy/api/jsonrpc.py
@@ -116,6 +116,7 @@ class Server():
             data = data.strip()
             logging.debug("Received %d bytes of data", len(data))
             self._handle_input(data)
+            await asyncio.sleep(0) # To not starve task queue
 
     def stop(self):
         self._active = False
diff --git a/src/galaxy/api/plugin.py b/src/galaxy/api/plugin.py
index 22eb9d4..43eb95f 100644
--- a/src/galaxy/api/plugin.py
+++ b/src/galaxy/api/plugin.py
@@ -38,6 +38,7 @@ class Plugin:
 
         self._feature_methods = OrderedDict()
         self._active = True
+        self._pass_control_task = None
 
         self._reader, self._writer = reader, writer
         self._handshake_token = handshake_token
@@ -210,15 +211,17 @@ class Plugin:
 
     async def run(self):
         """Plugin's main coroutine."""
-        async def pass_control():
-            while self._active:
-                try:
-                    self.tick()
-                except Exception:
-                    logging.exception("Unexpected exception raised in plugin tick")
-                await asyncio.sleep(1)
+        await self._server.run()
+        if self._pass_control_task is not None:
+            await self._pass_control_task
 
-        await asyncio.gather(pass_control(), self._server.run())
+    async def _pass_control(self):
+        while self._active:
+            try:
+                self.tick()
+            except Exception:
+                logging.exception("Unexpected exception raised in plugin tick")
+            await asyncio.sleep(1)
 
     def _shutdown(self):
         logging.info("Shutting down")
@@ -236,6 +239,7 @@ class Plugin:
     def _initialize_cache(self, data: Dict):
         self._persistent_cache = data
         self.handshake_complete()
+        self._pass_control_task = asyncio.create_task(self._pass_control())
 
     @staticmethod
     def _ping():
diff --git a/tests/test_internal.py b/tests/test_internal.py
index d381f4f..ec3dd77 100644
--- a/tests/test_internal.py
+++ b/tests/test_internal.py
@@ -62,7 +62,18 @@ def test_ping(plugin, read, write):
         "result": None
     }
 
-def test_tick(plugin, read):
+def test_tick_before_handshake(plugin, read):
     read.side_effect = [b""]
     asyncio.run(plugin.run())
+    plugin.tick.assert_not_called()
+
+def test_tick_after_handshake(plugin, read):
+    request = {
+        "jsonrpc": "2.0",
+        "id": "6",
+        "method": "initialize_cache",
+        "params": {"data": {}}
+    }
+    read.side_effect = [json.dumps(request).encode() + b"\n", b""]
+    asyncio.run(plugin.run())
     plugin.tick.assert_called_with()