Merge pull request #163 from plotly/neyberson/protocol_verify_json

Neyberson/protocol verify json

.github/workflows/test.yml

@@ -83,7 +83,7 @@ jobs:
       - name: Test Process Control
         if: ${{ ! runner.debug }}
         run:  pytest -W error -n auto -v -rFe --capture=fd tests/test_process.py
-        timeout-minutes: 2
+        timeout-minutes: 3
       - name: Test The Rest DEBUG
         if: runner.debug
         run:  pytest -W error -vvv -rA --capture=tee-sys --ignore=tests/test_process.py

choreographer/browser.py

@@ -523,6 +523,8 @@ async def close_tab(self, target_id):
             raise RuntimeError(
                 "There is no eventloop, or was not passed to browser. Cannot use async methods"
             )
+        if self.lock.locked():
+            raise BrowserClosedError("Calling commands after closed browser")
         if isinstance(target_id, Target):
             target_id = target_id.target_id
         # NOTE: we don't need to manually remove sessions because
@@ -531,6 +533,8 @@ async def close_tab(self, target_id):
             command="Target.closeTarget",
             params={"targetId": target_id},
         )
+        # TODO, without the lock, if we close and then call close_tab, does it hang like it did for
+        # test_tab_send_command in test_tab.py, or does it throw an error about a closed pipe?
         self._remove_tab(target_id)
         if "error" in response:
             raise RuntimeError("Could not close tab") from DevtoolsProtocolError(
@@ -623,11 +627,11 @@ def run_read_loop(self):
         def check_error(result):
             e = result.exception()
             if e:
-                if isinstance(e, asyncio.CancelledError):
-                    pass
-                elif self.debug:
-                    print(f"Error in run_read_loop: {str(e)}", file=sys.stderr)
                 self.close()
+                if self.debug:
+                    print(f"Error in run_read_loop: {str(e)}", file=sys.stderr)
+                if not isinstance(e, asyncio.CancelledError):
+                    raise e
         async def read_loop():
             try:
                 responses = await self.loop.run_in_executor(
@@ -701,6 +705,8 @@ async def read_loop():
                                     f"run_read_loop() found future for key {key}", file=sys.stderr
                                 )
                             future = self.futures.pop(key)
+                        elif "error" in response:
+                            raise DevtoolsProtocolError(response)
                         else:
                             raise RuntimeError(f"Couldn't find a future for key: {key}")
                         if error:

choreographer/protocol.py

@@ -8,6 +8,21 @@ def __init__(self, response):
         self.message = response["error"]["message"]
 
 
+class MessageTypeError(TypeError):
+    def __init__(self, key, value, expected_type):
+        value = type(value) if not isinstance(value, type) else value
+        super().__init__(
+            f"Message with key {key} must have type {expected_type}, not {value}."
+        )
+
+
+class MissingKeyError(ValueError):
+    def __init__(self, key, obj):
+        super().__init__(
+            f"Message missing required key/s {key}. Message received: {obj}"
+        )
+
+
 class ExperimentalFeatureWarning(UserWarning):
     pass
 
@@ -31,10 +46,13 @@ def calculate_key(self, response):
 
     def verify_json(self, obj):
         n_keys = 0
-        if "id" in obj and "method" in obj:
-            n_keys += 2
-        else:
-            raise RuntimeError("Each message object must contain an id and method key")
+        required_keys = {"id": int, "method": str}
+        for key, type_key in required_keys.items():
+            if key not in obj:
+                raise MissingKeyError(key, obj)
+            if not isinstance(obj[key], type_key):
+                raise MessageTypeError(key, type(obj[key]), type_key)
+        n_keys += 2
 
         if "params" in obj:
             n_keys += 1

tests/conftest.py

@@ -48,7 +48,10 @@ async def browser(request):
         headless=headless, debug=debug, debug_browser=debug
     )
     yield browser
-    await browser.close()
+    try:
+        await browser.close()
+    except choreo.browser.BrowserClosedError:
+        pass
 
 
 

tests/test_browser.py

@@ -30,15 +30,69 @@ async def test_create_and_close_session(browser):
 
 @pytest.mark.asyncio
 async def test_browser_write_json(browser):
+    # Test valid request with correct id and method
     response = await browser.write_json({"id": 0, "method": "Target.getTargets"})
     assert "result" in response and "targetInfos" in response["result"]
 
+    # Test invalid method name should return error
+    response = await browser.write_json({"id": 2, "method": "dkadklqwmd"})
+    assert "error" in response
+
+    # Test missing 'id' key
+    with pytest.raises(
+        choreo.protocol.MissingKeyError,
+    ):
+        await browser.write_json({"method": "Target.getTargets"})
+
+    # Test missing 'method' key
+    with pytest.raises(
+        choreo.protocol.MissingKeyError,
+    ):
+        await browser.write_json({"id": 1})
+
+    # Test empty dictionary
+    with pytest.raises(
+        choreo.protocol.MissingKeyError,
+    ):
+        await browser.write_json({})
+
+    # Test invalid parameter in the message
+    with pytest.raises(
+        RuntimeError,
+    ):
+        await browser.write_json(
+            {"id": 0, "method": "Target.getTargets", "invalid_parameter": "kamksamdk"}
+        )
+
+    # Test int method should return error
+    with pytest.raises(
+        choreo.protocol.MessageTypeError,
+    ):
+        await browser.write_json({"id": 3, "method": 12345})
+
+    # Test non-integer id should return error
+    with pytest.raises(
+        choreo.protocol.MessageTypeError,
+    ):
+        await browser.write_json({"id": "2", "method": "Target.getTargets"})
+
 
 @pytest.mark.asyncio
 async def test_browser_send_command(browser):
+    # Test valid request with correct command
     response = await browser.send_command(command="Target.getTargets")
     assert "result" in response and "targetInfos" in response["result"]
 
+    # Test invalid method name should return error
+    response = await browser.send_command(command="dkadklqwmd")
+    assert "error" in response
+
+    # Test int method should return error
+    with pytest.raises(
+        choreo.protocol.MessageTypeError,
+    ):
+        await browser.send_command(command=12345)
+
 
 @pytest.mark.asyncio
 async def test_populate_targets(browser):

tests/test_session.py

@@ -16,6 +16,17 @@ async def session(browser):
 
 
 @pytest.mark.asyncio
-async def test_send_command(session):
-    response = await session.send_command("Target.getTargets")
+async def test_session_send_command(session):
+    # Test valid request with correct command
+    response = await session.send_command(command="Target.getTargets")
     assert "result" in response and "targetInfos" in response["result"]
+
+    # Test invalid method name should return error
+    response = await session.send_command(command="dkadklqwmd")
+    assert "error" in response
+
+    # Test int method should return error
+    with pytest.raises(
+        choreo.protocol.MessageTypeError,
+    ):
+        await session.send_command(command=12345)

tests/test_tab.py

@@ -10,14 +10,19 @@ def check_response_dictionary(response_received, response_expected):
     for k, v in response_expected.items():
         if isinstance(v, dict):
             check_response_dictionary(v, response_expected[k])
-        assert k in response_received and response_received[k] == v, "Expected: {response_expected}\nReceived: {response_received}"
+        assert (
+            k in response_received and response_received[k] == v
+        ), "Expected: {response_expected}\nReceived: {response_received}"
 
 
 @pytest_asyncio.fixture(scope="function", loop_scope="function")
 async def tab(browser):
     tab_browser = await browser.create_tab("")
     yield tab_browser
-    await browser.close_tab(tab_browser)
+    try:
+        await browser.close_tab(tab_browser)
+    except choreo.browser.BrowserClosedError:
+        pass
 
 
 @pytest.mark.asyncio
@@ -29,10 +34,21 @@ async def test_create_and_close_session(tab):
 
 
 @pytest.mark.asyncio
-async def test_send_command(tab):
-    response = await tab.send_command("Page.enable")
+async def test_tab_send_command(tab):
+    # Test valid request with correct command
+    response = await tab.send_command(command="Page.enable")
     check_response_dictionary(response, {"result": {}})
 
+    # Test invalid method name should return error
+    response = await tab.send_command(command="dkadklqwmd")
+    assert "error" in response
+
+    # Test int method should return error
+    with pytest.raises(
+        choreo.protocol.MessageTypeError,
+    ):
+        await tab.send_command(command=12345)
+
 
 @pytest.mark.asyncio
 async def test_subscribe_once(tab):
@@ -47,15 +63,18 @@ async def test_subscribe_once(tab):
 @pytest.mark.asyncio
 async def test_subscribe_and_unsubscribe(tab):
     counter = 0
+    old_counter = counter
+
     async def count_event(r):
         nonlocal counter
         counter += 1
+
     tab.subscribe("Page.*", count_event)
     assert "Page.*" in list(tab.sessions.values())[0].subscriptions
     await tab.send_command("Page.enable")
     await tab.send_command("Page.reload")
-    await asyncio.sleep(.5)
-    assert counter > 1
+    await asyncio.sleep(0.5)
+    assert counter > old_counter
 
     tab.unsubscribe("Page.*")
     old_counter = counter