Test more of recent AgentKey additions

2 files changed, 25 insertions, 65 deletions

diff --git a/tests/agent.py b/tests/agent.py
index 18ece7a9..158a8b41 100644
--- a/tests/agent.py
+++ b/tests/agent.py
@@ -1,24 +1,41 @@
 from paramiko import AgentKey
 
 
+# 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_:
+    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 fields:
         """
         _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"]
 
         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()
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)