diff --git a/hippolyzer/lib/proxy/commands.py b/hippolyzer/lib/proxy/commands.py index 6662225..d32de70 100644 --- a/hippolyzer/lib/proxy/commands.py +++ b/hippolyzer/lib/proxy/commands.py @@ -65,13 +65,13 @@ def handle_command(command_name: Optional[str] = None, /, *, lifetime: Optional[ # Greedy, takes the rest of the message if param.sep is None: param_val = message - message = None + message = "" else: message = message.lstrip(param.sep) if not message: - if param.optional: - break - raise KeyError(f"Missing parameter {param_name}") + if not param.optional: + raise KeyError(f"Missing parameter {param_name}") + continue param_val, _, message = message.partition(param.sep) # type: ignore param_vals[param_name] = param.parser(param_val) diff --git a/tests/proxy/test_commands.py b/tests/proxy/test_commands.py index b6c2d90..e5914c5 100644 --- a/tests/proxy/test_commands.py +++ b/tests/proxy/test_commands.py @@ -26,7 +26,13 @@ class ExampleCommandHandler: y=str, ) async def own_name(self, _session, _region, y): - self.bar = y + pass + + @handle_command( + x=Parameter(str, optional=True), + ) + async def optional(self, _session, _region, x=42): + self.bar = x class TestCommandHandlers(unittest.IsolatedAsyncioTestCase): @@ -47,6 +53,17 @@ class TestCommandHandlers(unittest.IsolatedAsyncioTestCase): async def test_own_name(self): self.assertEqual(self.handler.own_name.command.name, "own_name") + async def test_missing_param(self): + with self.assertRaises(KeyError): + await self.handler.foo(None, None, "") + + async def test_optional_param(self): + await self.handler.optional(None, None, "foo") # type: ignore + self.assertEqual(self.handler.bar, "foo") + await self.handler.optional(None, None, "") # type: ignore + # Should have picked up the default value + self.assertEqual(self.handler.bar, 42) + async def test_bad_command(self): with self.assertRaises(ValueError): class _BadCommandHandler: