Test more of recent AgentKey additions

Migrate rest of existing agent tests to new module

Also make them better.

Also fix AttributeError raising to not be py3.10 specific

2 files changed, 95 insertions, 70 deletions

diff --git a/tests/agent.py b/tests/agent.py
index 18ece7a9..8e859289 100644
--- a/tests/agent.py
+++ b/tests/agent.py
@@ -1,24 +1,106 @@
-from paramiko import AgentKey
+from unittest.mock import Mock
+
+from pytest import mark, raises
+
+from paramiko import AgentKey, Message
+from paramiko.agent import (
+    SSH2_AGENT_SIGN_RESPONSE,
+    SSH_AGENT_RSA_SHA2_256,
+    SSH_AGENT_RSA_SHA2_512,
+    cSSH2_AGENTC_SIGN_REQUEST,
+)
+
+
+# AgentKey with no inner_key
+class _BareAgentKey(AgentKey):
+    def __init__(self, name, blob):
+        self.name = name
+        self.blob = blob
+        self.inner_key = None
 
 
 class AgentKey_:
-    class fields:
-        """
-        _fields
-        """
+    def str_is_repr(self):
+        # Tests for a missed spot in Python 3 upgrades: AgentKey.__str__ was
+        # returning bytes, as if under Python 2. When bug present, this
+        # explodes with "__str__ returned non-string".
+        key = AgentKey(None, b"secret!!!")
+        assert str(key) == repr(key)
+
+    class init:
+        def needs_at_least_two_arguments(self):
+            with raises(TypeError):
+                AgentKey()
+            with raises(TypeError):
+                AgentKey(None)
+
+        def sets_attributes_and_parses_blob(self):
+            agent = Mock()
+            blob = Message()
+            blob.add_string("bad-type")
+            key = AgentKey(agent=agent, blob=bytes(blob))
+            assert key.agent is agent
+            assert key.name == "bad-type"
+            assert key.blob == bytes(blob)
+            assert key.comment == ""  # default
+            # TODO: logger testing
+            assert key.inner_key is None  # no 'bad-type' algorithm
+
+        def comment_optional(self):
+            blob = Message()
+            blob.add_string("bad-type")
+            key = AgentKey(agent=Mock(), blob=bytes(blob), comment="hi!")
+            assert key.comment == "hi!"
 
+        def sets_inner_key_when_known_type(self, keys):
+            key = AgentKey(agent=Mock(), blob=bytes(keys.pkey))
+            assert key.inner_key == keys.pkey
+
+    class fields:
         def defaults_to_get_name_and_blob(self):
-            # Manually construct a 'failed to get inner_key' obj
-            class FallbackAgentKey(AgentKey):
-                def __init__(self, name, blob):
-                    self.name = name
-                    self.blob = blob
-                    self.inner_key = None
-
-            key = FallbackAgentKey(name="lol", blob=b"lmao")
+            key = _BareAgentKey(name="lol", blob=b"lmao")
             assert key._fields == ["lol", b"lmao"]
 
+        # TODO: pytest-relaxed is buggy (now?), this shows up under get_bits?
         def defers_to_inner_key_when_present(self, keys):
             key = AgentKey(agent=None, blob=keys.pkey.asbytes())
             assert key._fields == keys.pkey._fields
             assert key == keys.pkey
+
+    class get_bits:
+        def defaults_to_superclass_implementation(self):
+            # TODO 4.0: assert raises NotImplementedError like changed parent?
+            assert _BareAgentKey(None, None).get_bits() == 0
+
+        def defers_to_inner_key_when_present(self, keys):
+            key = AgentKey(agent=None, blob=keys.pkey.asbytes())
+            assert key.get_bits() == keys.pkey.get_bits()
+
+    @mark.parametrize(
+        "kwargs,expectation",
+        [
+            # No algorithm kwarg: no flags (bitfield -> 0 int)
+            (dict(), 0),
+            (dict(algorithm="rsa-sha2-256"), SSH_AGENT_RSA_SHA2_256),
+            (dict(algorithm="rsa-sha2-512"), SSH_AGENT_RSA_SHA2_512),
+        ],
+    )
+    def signing_data(self, kwargs, expectation):
+        class FakeAgent:
+            def _send_message(self, msg):
+                self._sent_message = msg
+                sig = Message()
+                sig.add_string("lol")
+                sig.rewind()
+                return SSH2_AGENT_SIGN_RESPONSE, sig
+
+        agent = FakeAgent()
+        key = AgentKey(agent, b"secret!!!")
+        result = key.sign_ssh_data(b"token", **kwargs)
+        assert result == b"lol"
+        msg = agent._sent_message
+        msg.rewind()
+        assert msg.get_byte() == cSSH2_AGENTC_SIGN_REQUEST
+        assert msg.get_string() == b"secret!!!"
+        assert msg.get_string() == b"token"
+        assert msg.get_int() == expectation
diff --git a/tests/test_agent.py b/tests/test_agent.py
deleted file mode 100644
index a00011bf..00000000
--- a/tests/test_agent.py
+++ /dev/null
@@ -1,57 +0,0 @@
-import unittest
-
-from paramiko.message import Message
-from paramiko.agent import (
-    SSH2_AGENT_SIGN_RESPONSE,
-    cSSH2_AGENTC_SIGN_REQUEST,
-    SSH_AGENT_RSA_SHA2_256,
-    SSH_AGENT_RSA_SHA2_512,
-    AgentKey,
-)
-from paramiko.util import b
-
-
-class ChaosAgent:
-    def _send_message(self, msg):
-        self._sent_message = msg
-        sig = Message()
-        sig.add_string(b("lol"))
-        sig.rewind()
-        return SSH2_AGENT_SIGN_RESPONSE, sig
-
-
-class AgentTests(unittest.TestCase):
-    def _sign_with_agent(self, kwargs, expectation):
-        agent = ChaosAgent()
-        key = AgentKey(agent, b("secret!!!"))
-        result = key.sign_ssh_data(b("token"), **kwargs)
-        assert result == b("lol")
-        msg = agent._sent_message
-        msg.rewind()
-        assert msg.get_byte() == cSSH2_AGENTC_SIGN_REQUEST
-        assert msg.get_string() == b("secret!!!")
-        assert msg.get_string() == b("token")
-        assert msg.get_int() == expectation
-
-    def test_agent_signing_defaults_to_0_for_flags_field(self):
-        # No algorithm kwarg at all
-        self._sign_with_agent(kwargs=dict(), expectation=0)
-
-    def test_agent_signing_is_2_for_SHA256(self):
-        self._sign_with_agent(
-            kwargs=dict(algorithm="rsa-sha2-256"),
-            expectation=SSH_AGENT_RSA_SHA2_256,
-        )
-
-    def test_agent_signing_is_2_for_SHA512(self):
-        self._sign_with_agent(
-            kwargs=dict(algorithm="rsa-sha2-512"),
-            expectation=SSH_AGENT_RSA_SHA2_512,
-        )
-
-    def test_agent_key_str_kinda_fixed(self):
-        # Tests for a missed spot in Python 3 upgrades: AgentKey.__str__ was
-        # returning bytes, as if under Python 2. When bug present, this
-        # explodes with "__str__ returned non-string".
-        key = AgentKey(ChaosAgent(), b("secret!!!"))
-        assert str(key) == repr(key)