Refactor command handling

This now supports usage/help texts

tests/integration/test_sshd.py

@@ -69,8 +69,8 @@ class TestSshd:
             assert found is True
             session.stdin.write(test_client.public_key + "\n")
 
-            result = await session.stdout.readline()
-            assert "OK" in result
+            result = await session.stdout.read()
+            assert "Key is valid. Registering client." in result
             await session.wait()
         return test_client
 

tests/packages/sshd/conftest.py

@@ -82,10 +82,13 @@ async def mock_backend(client_registry: ClientRegistry) -> MagicMock:
                 "Error, must have a client called template for this to work."
             )
         clients_data[name] = clients_data["template"]
+        template_secrets: dict[str, str] = {}
         for secret_key, secret in secrets_data.items():
             s_client, secret_name = secret_key
             if s_client != "template":
                 continue
+            template_secrets[secret_name] = secret
+        for secret_name, secret in template_secrets.items():
             secrets_data[(name, secret_name)] = secret
 
     async def write_audit(*args, **kwargs):

tests/packages/sshd/test_errors.py

@@ -83,8 +83,9 @@ class TestRegistrationErrors(BaseSshTests):
                 output = await process.stdout.readline()
                 assert "Enter public key" in output
                 stdout, stderr = await process.communicate(public_key)
+                assert isinstance(stderr, str)
                 print(f"{stdout=!r}, {stderr=!r}")
-                assert stderr == "Error: Invalid key type: Only RSA keys are supported."
+                assert stderr.rstrip() == "Error: Invalid key type: Only RSA keys are supported."
                 result = await process.wait()
                 assert result.exit_status == 1
 
@@ -102,8 +103,9 @@ class TestRegistrationErrors(BaseSshTests):
                 output = await process.stdout.readline()
                 assert "Enter public key" in output
                 stdout, stderr = await process.communicate(public_key)
+                assert isinstance(stderr, str)
                 print(f"{stdout=!r}, {stderr=!r}")
-                assert stderr == "Error: Invalid key type: Only RSA keys are supported."
+                assert stderr.rstrip() == "Error: Invalid key type: Only RSA keys are supported."
                 result = await process.wait()
                 assert result.exit_status == 1
 
@@ -122,7 +124,8 @@ class TestCommandErrors(BaseSshTests):
 
         assert result.exit_status == 1
         stderr = result.stderr or ""
-        assert stderr == "Error: Unsupported command."
+        assert isinstance(stderr, str)
+        assert stderr.rstrip() == "Error: Unsupported command."
 
     @pytest.mark.asyncio
     async def test_no_command(
@@ -136,7 +139,8 @@ class TestCommandErrors(BaseSshTests):
             async with conn.create_process() as process:
                 stdout, stderr = await process.communicate()
                 print(f"{stdout=!r}, {stderr=!r}")
-                assert stderr == "Error: No command was received from the client."
+                assert isinstance(stderr, str)
+                assert stderr.rstrip() == "Error: No command was received from the client."
                 result = await process.wait()
                 assert result.exit_status == 1
 

tests/packages/sshd/test_get_secret.py

@@ -1,10 +1,12 @@
 """Test get secret."""
 
+import allure
 import pytest
 
 from .types import ClientRegistry, CommandRunner
 
 
+@allure.title("Test get_secret command")
 @pytest.mark.asyncio
 async def test_get_secret(
     ssh_command_runner: CommandRunner, client_registry: ClientRegistry
@@ -19,7 +21,7 @@ async def test_get_secret(
     assert isinstance(result.stdout, str)
     assert result.stdout.rstrip() == "mocked-secret-mysecret"
 
-
+@allure.title("Test with invalid secret name")
 @pytest.mark.asyncio
 async def test_invalid_secret_name(
     ssh_command_runner: CommandRunner, client_registry: ClientRegistry
@@ -30,4 +32,25 @@ async def test_invalid_secret_name(
 
     result = await ssh_command_runner("test-client", "get_secret mysecret")
     assert result.exit_status == 1
-    assert result.stderr == "Error: No secret available with the given name."
+    stderr = result.stderr
+    assert isinstance(stderr, str)
+    assert stderr.rstrip() == "Error: No secret available with the given name."
+
+@allure.title("Test get_secret command help")
+@pytest.mark.asyncio
+async def test_get_secret_cmd_help(ssh_command_runner: CommandRunner,  client_registry: ClientRegistry) -> None:
+    """Test running get_secret --help"""
+    await client_registry["add_client"]("test-client", ["mysecret"])
+
+    result = await ssh_command_runner("test-client", "get_secret --help")
+
+    assert result.exit_status == 0
+
+    print(result.stdout)
+    assert isinstance(result.stdout, str)
+
+    lines = result.stdout.splitlines()
+
+    assert lines[0] == "get_secret SECRET"
+
+    assert len(lines) > 4

tests/packages/sshd/test_ping.py

@@ -1,8 +1,11 @@
+"""Test for the ping command."""
+import allure
 import pytest
 
 from .types import ClientRegistry, CommandRunner
 
 
+@allure.title("Test running the ping command")
 @pytest.mark.asyncio
 async def test_ping_command(
     ssh_command_runner: CommandRunner, client_registry: ClientRegistry
@@ -16,3 +19,21 @@ async def test_ping_command(
     assert result.stdout is not None
     assert isinstance(result.stdout, str)
     assert result.stdout.rstrip() == "PONG"
+
+@allure.title("Test ping help")
+@pytest.mark.asyncio
+async def test_ping_cmd_help(ssh_command_runner: CommandRunner,  client_registry: ClientRegistry) -> None:
+    """Test running ping --help."""
+    await client_registry["add_client"]("test-client", ["mysecret"])
+    result = await ssh_command_runner("test-client", "ping --help")
+
+    assert result.exit_status == 0
+
+    print(result.stdout)
+    assert isinstance(result.stdout, str)
+
+    lines = result.stdout.splitlines()
+
+    assert lines[0] == "ping"
+
+    assert len(lines) > 4

tests/packages/sshd/test_register.py

@@ -1,10 +1,12 @@
 """Test registration."""
 
+import allure
 import pytest
 
 from .types import ClientRegistry, CommandRunner, ProcessRunner
 
 
+@allure.title("Test client registration")
 @pytest.mark.enable_registration(True)
 @pytest.mark.asyncio
 async def test_register_client(
@@ -29,8 +31,9 @@ async def test_register_client(
 
         assert found is True
         session.stdin.write(public_key)
-        result = await session.stdout.readline()
-        assert "OK" in result
+        data = await session.stdout.read()
+        assert isinstance(data, str)
+        assert "Key is valid. Registering client" in data
 
     # Test that we can connect
 
@@ -39,3 +42,28 @@ async def test_register_client(
     assert result.stdout is not None
     assert isinstance(result.stdout, str)
     assert result.stdout.rstrip() == "mocked-secret-testsecret"
+
+@allure.title("Test register command help")
+@pytest.mark.enable_registration(True)
+@pytest.mark.asyncio
+async def test_register_cmd_help(ssh_command_runner: CommandRunner,  client_registry: ClientRegistry) -> None:
+    """Test running register --help"""
+    await client_registry["add_client"]("test-client", ["mysecret"])
+
+    result = await ssh_command_runner("test-client", "register --help")
+
+    assert result.exit_status == 0
+
+    print(result.stdout)
+    assert isinstance(result.stdout, str)
+
+    lines = result.stdout.splitlines()
+
+    assert lines[0] == "register"
+
+    assert len(lines) > 4
+
+
+
+
+# TODO: Test running register with an existing client.