gabriel / musehub public
test_agent_registration.py python
460 lines 16.2 KB
Raw
sha256:9590cee1e0ccd6c76528f005b95d634d80f5019f0dcb7c371e149adc31d1fb65 refactor: enforce gRPC framing on all MWP wire traffic Sonnet 4.6 minor ⚠ breaking 156 days ago
1 """Tests for agent identity provisioning.
2
3 Covers:
4 - AgentRegistrationRequest model validation
5 - register_agent_identity service function (unit tests with mocked DB)
6 - POST /api/identities/agent route (happy path + error cases)
7 - verify_and_authenticate identity_type support
8 - VerifyRequest identity_type field validation
9 """
10 from __future__ import annotations
11
12 import base64
13 import hashlib
14 import uuid
15 from unittest.mock import AsyncMock, MagicMock, patch
16
17 import pytest
18 from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
19 from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
20 from httpx import AsyncClient
21
22 from musehub.types.json_types import StrDict
23 from musehub.models.musehub_auth import (
24 AgentRegistrationRequest,
25 AgentRegistrationResponse,
26 VerifyRequest,
27 )
28
29
30 # ---------------------------------------------------------------------------
31 # Helpers
32 # ---------------------------------------------------------------------------
33
34
35 def _generate_key_material() -> tuple[str, str]:
36 """Generate a fresh Ed25519 keypair and return (public_key_b64, fingerprint)."""
37 key = Ed25519PrivateKey.generate()
38 pub_raw = key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
39 pub_b64 = base64.urlsafe_b64encode(pub_raw).rstrip(b"=").decode("ascii")
40 fingerprint = hashlib.sha256(pub_raw).hexdigest()
41 return pub_b64, fingerprint
42
43
44 # ---------------------------------------------------------------------------
45 # AgentRegistrationRequest validation
46 # ---------------------------------------------------------------------------
47
48
49 class TestAgentRegistrationRequestValidation:
50 def _valid_payload(self) -> JSONObject:
51 pub_b64, fp = _generate_key_material()
52 return {
53 "handle": "agentception-abc123",
54 "public_key_b64": pub_b64,
55 "fingerprint": fp,
56 "algorithm": "ed25519",
57 "agent_model": "claude-sonnet-4-6",
58 "scope": ["push:agentception"],
59 "label": "ephemeral/agentception-abc123",
60 }
61
62 def test_valid_request_parses(self) -> None:
63 req = AgentRegistrationRequest(**self._valid_payload())
64 assert req.handle == "agentception-abc123"
65 assert req.algorithm == "ed25519"
66 assert req.agent_model == "claude-sonnet-4-6"
67 assert req.scope == ["push:agentception"]
68
69 def test_handle_normalised_to_lowercase(self) -> None:
70 payload = self._valid_payload()
71 payload["handle"] = "AgentCeption-ABC"
72 req = AgentRegistrationRequest(**payload)
73 assert req.handle == "agentception-abc"
74
75 def test_invalid_handle_rejected(self) -> None:
76 payload = self._valid_payload()
77 payload["handle"] = "agent with spaces"
78 with pytest.raises(Exception): # ValidationError
79 AgentRegistrationRequest(**payload)
80
81 def test_fingerprint_must_be_64_hex(self) -> None:
82 payload = self._valid_payload()
83 payload["fingerprint"] = "tooshort"
84 with pytest.raises(Exception):
85 AgentRegistrationRequest(**payload)
86
87 def test_invalid_algorithm_rejected(self) -> None:
88 payload = self._valid_payload()
89 payload["algorithm"] = "rsa-2048"
90 with pytest.raises(Exception):
91 AgentRegistrationRequest(**payload)
92
93 def test_scope_defaults_to_empty_list(self) -> None:
94 payload = self._valid_payload()
95 del payload["scope"]
96 req = AgentRegistrationRequest(**payload)
97 assert req.scope == []
98
99 def test_expires_at_defaults_to_none(self) -> None:
100 req = AgentRegistrationRequest(**self._valid_payload())
101 assert req.expires_at is None
102
103 def test_expires_at_accepts_iso8601(self) -> None:
104 payload = self._valid_payload()
105 payload["expires_at"] = "2026-04-06T14:00:00Z"
106 req = AgentRegistrationRequest(**payload)
107 assert req.expires_at == "2026-04-06T14:00:00Z"
108
109
110 # ---------------------------------------------------------------------------
111 # VerifyRequest identity_type field
112 # ---------------------------------------------------------------------------
113
114
115 class TestVerifyRequestIdentityType:
116 def _base_payload(self) -> JSONObject:
117 return {
118 "challenge_token": "a" * 64,
119 "public_key_b64": "AAEC",
120 "signature_b64": "AAEC",
121 }
122
123 def test_default_identity_type_is_human(self) -> None:
124 req = VerifyRequest(**self._base_payload())
125 assert req.identity_type == "human"
126
127 def test_agent_identity_type_accepted(self) -> None:
128 req = VerifyRequest(**self._base_payload(), identity_type="agent")
129 assert req.identity_type == "agent"
130
131 def test_invalid_identity_type_rejected(self) -> None:
132 with pytest.raises(Exception):
133 VerifyRequest(**self._base_payload(), identity_type="robot")
134
135
136 # ---------------------------------------------------------------------------
137 # register_agent_identity service — unit tests
138 # ---------------------------------------------------------------------------
139
140
141 class TestRegisterAgentIdentityService:
142 """Unit tests using a mock DB session — no real DB required."""
143
144 @staticmethod
145 def _make_created_at():
146 from datetime import datetime, timezone
147 t = MagicMock()
148 t.isoformat.return_value = "2026-04-06T00:00:00+00:00"
149 return t
150
151 def _make_mock_session(
152 self, *, key_row_exists: bool = False, identity_row_exists: bool = False
153 ):
154 session = AsyncMock()
155
156 # Simulated scalar_one_or_none return for SELECT MusehubAuthKey
157 mock_scalar_result = MagicMock()
158 if key_row_exists:
159 mock_key = MagicMock()
160 mock_key.key_id = str(uuid.uuid4())
161 mock_key.identity_id = str(uuid.uuid4())
162 mock_key.algorithm = "ed25519"
163 mock_key.fingerprint = "a" * 64
164 mock_key.label = "existing"
165 mock_key.created_at = self._make_created_at()
166 mock_key.last_used_at = None
167 mock_scalar_result.scalar_one_or_none.return_value = mock_key
168
169 if identity_row_exists:
170 mock_identity = MagicMock()
171 mock_identity.identity_id = mock_key.identity_id
172 mock_identity.handle = "agentception-abc"
173 # second execute call returns identity
174 mock_scalar_result2 = MagicMock()
175 mock_scalar_result2.scalar_one_or_none.return_value = mock_identity
176 session.execute.side_effect = [
177 mock_scalar_result,
178 mock_scalar_result2,
179 ]
180 else:
181 session.execute.return_value = mock_scalar_result
182 else:
183 mock_scalar_result.scalar_one_or_none.return_value = None
184 session.execute.return_value = mock_scalar_result
185
186 session.flush = AsyncMock()
187 session.commit = AsyncMock()
188 session.add = MagicMock()
189
190 # refresh populates created_at on any ORM object passed to it
191 async def _mock_refresh(obj):
192 if not hasattr(obj, "created_at") or obj.created_at is None:
193 obj.created_at = self._make_created_at()
194 if not hasattr(obj, "last_used_at"):
195 obj.last_used_at = None
196
197 session.refresh = _mock_refresh
198
199 return session
200
201 @pytest.mark.asyncio
202 async def test_new_agent_registration_creates_identity_and_key(self) -> None:
203 from musehub.services.musehub_auth import register_agent_identity
204
205 pub_b64, fp = _generate_key_material()
206 session = self._make_mock_session(key_row_exists=False)
207
208 result = await register_agent_identity(
209 session=session,
210 handle="agentception-abc",
211 public_key_b64=pub_b64,
212 fingerprint=fp,
213 algorithm="ed25519",
214 spawned_by="gabriel",
215 agent_model="claude-sonnet-4-6",
216 scope=["push:agentception"],
217 )
218
219 assert result.is_new_identity is True
220 assert result.spawned_by == "gabriel"
221 assert result.handle == "agentception-abc"
222 # session.add called twice: identity + key
223 assert session.add.call_count == 2
224 assert session.commit.call_count == 1
225
226 @pytest.mark.asyncio
227 async def test_fingerprint_mismatch_raises_auth_error(self) -> None:
228 from musehub.services.musehub_auth import AuthError, register_agent_identity
229
230 pub_b64, _ = _generate_key_material()
231 wrong_fp = "f" * 64 # doesn't match the key
232
233 session = self._make_mock_session(key_row_exists=False)
234
235 with pytest.raises(AuthError) as exc_info:
236 await register_agent_identity(
237 session=session,
238 handle="agent-x",
239 public_key_b64=pub_b64,
240 fingerprint=wrong_fp,
241 algorithm="ed25519",
242 spawned_by="gabriel",
243 )
244
245 assert exc_info.value.status_code == 422
246 assert "fingerprint" in exc_info.value.detail.lower()
247
248 @pytest.mark.asyncio
249 async def test_invalid_public_key_b64_raises_auth_error(self) -> None:
250 from musehub.services.musehub_auth import AuthError, register_agent_identity
251
252 session = self._make_mock_session(key_row_exists=False)
253
254 with pytest.raises(AuthError) as exc_info:
255 await register_agent_identity(
256 session=session,
257 handle="agent-x",
258 public_key_b64="!!!not-base64!!!",
259 fingerprint="a" * 64,
260 algorithm="ed25519",
261 spawned_by="gabriel",
262 )
263
264 assert exc_info.value.status_code == 422
265
266 @pytest.mark.asyncio
267 async def test_invalid_expires_at_raises_auth_error(self) -> None:
268 from musehub.services.musehub_auth import AuthError, register_agent_identity
269
270 pub_b64, fp = _generate_key_material()
271 session = self._make_mock_session(key_row_exists=False)
272
273 with pytest.raises(AuthError) as exc_info:
274 await register_agent_identity(
275 session=session,
276 handle="agent-x",
277 public_key_b64=pub_b64,
278 fingerprint=fp,
279 algorithm="ed25519",
280 spawned_by="gabriel",
281 expires_at="not-a-date",
282 )
283
284 assert exc_info.value.status_code == 422
285 assert "expires_at" in exc_info.value.detail.lower()
286
287 @pytest.mark.asyncio
288 async def test_expires_at_none_is_accepted(self) -> None:
289 from musehub.services.musehub_auth import register_agent_identity
290
291 pub_b64, fp = _generate_key_material()
292 session = self._make_mock_session(key_row_exists=False)
293
294 result = await register_agent_identity(
295 session=session,
296 handle="agent-no-expiry",
297 public_key_b64=pub_b64,
298 fingerprint=fp,
299 algorithm="ed25519",
300 spawned_by="gabriel",
301 expires_at=None,
302 )
303 assert result.is_new_identity is True
304
305
306 # ---------------------------------------------------------------------------
307 # POST /api/identities/agent route — HTTP integration tests
308 # ---------------------------------------------------------------------------
309
310
311 class TestProvisionAgentRoute:
312 """Integration tests using the full FastAPI test client."""
313
314 def _valid_payload(self) -> JSONObject:
315 pub_b64, fp = _generate_key_material()
316 return {
317 "handle": f"agent-{uuid.uuid4().hex[:8]}",
318 "public_key_b64": pub_b64,
319 "fingerprint": fp,
320 "algorithm": "ed25519",
321 "agent_model": "claude-sonnet-4-6",
322 "scope": ["push"],
323 "label": "test-ephemeral",
324 }
325
326 @pytest.mark.asyncio
327 async def test_provision_agent_happy_path(
328 self,
329 client: AsyncClient,
330 auth_headers: StrDict,
331 db_session,
332 ) -> None:
333 payload = self._valid_payload()
334 resp = await client.post(
335 "/api/identities/agent",
336 json=payload,
337 headers=auth_headers,
338 )
339 assert resp.status_code in (200, 201)
340 data = resp.json()
341 assert data["handle"] == payload["handle"]
342 assert data["spawned_by"] == "testuser" # from auth_headers fixture
343 assert data["is_new_identity"] is True
344 assert "key" in data
345 assert data["key"]["algorithm"] == "ed25519"
346
347 @pytest.mark.asyncio
348 async def test_provision_agent_requires_auth(
349 self,
350 client: AsyncClient,
351 ) -> None:
352 payload = self._valid_payload()
353 resp = await client.post(
354 "/api/identities/agent",
355 json=payload,
356 # No auth headers
357 )
358 assert resp.status_code == 401
359
360 @pytest.mark.asyncio
361 async def test_provision_agent_duplicate_handle_409(
362 self,
363 client: AsyncClient,
364 auth_headers: StrDict,
365 db_session,
366 ) -> None:
367 payload = self._valid_payload()
368 # First registration
369 r1 = await client.post("/api/identities/agent", json=payload, headers=auth_headers)
370 assert r1.status_code in (200, 201)
371
372 # Same handle with a different key — should 409
373 _, fp2 = _generate_key_material()
374 pub_b64_2, fp2 = _generate_key_material()
375 payload2 = {**payload, "public_key_b64": pub_b64_2, "fingerprint": fp2}
376 r2 = await client.post("/api/identities/agent", json=payload2, headers=auth_headers)
377 assert r2.status_code == 409
378
379 @pytest.mark.asyncio
380 async def test_provision_agent_idempotent_same_key(
381 self,
382 client: AsyncClient,
383 auth_headers: StrDict,
384 db_session,
385 ) -> None:
386 """Registering the same key twice is idempotent (returns 200 on re-register)."""
387 payload = self._valid_payload()
388 r1 = await client.post("/api/identities/agent", json=payload, headers=auth_headers)
389 assert r1.status_code in (200, 201)
390
391 r2 = await client.post("/api/identities/agent", json=payload, headers=auth_headers)
392 assert r2.status_code == 200
393 data2 = r2.json()
394 assert data2["is_new_identity"] is False
395
396 @pytest.mark.asyncio
397 async def test_provision_agent_invalid_handle_422(
398 self,
399 client: AsyncClient,
400 auth_headers: StrDict,
401 ) -> None:
402 payload = self._valid_payload()
403 payload["handle"] = "handle with spaces"
404 resp = await client.post("/api/identities/agent", json=payload, headers=auth_headers)
405 assert resp.status_code == 422
406
407 @pytest.mark.asyncio
408 async def test_provision_agent_spawned_by_matches_operator(
409 self,
410 client: AsyncClient,
411 auth_headers: StrDict,
412 db_session,
413 ) -> None:
414 payload = self._valid_payload()
415 resp = await client.post("/api/identities/agent", json=payload, headers=auth_headers)
416 assert resp.status_code in (200, 201)
417 data = resp.json()
418 # _TEST_HANDLE from conftest
419 assert data["spawned_by"] == "testuser"
420
421 @pytest.mark.asyncio
422 async def test_provision_agent_scope_stored(
423 self,
424 client: AsyncClient,
425 auth_headers: StrDict,
426 db_session,
427 ) -> None:
428 payload = self._valid_payload()
429 payload["scope"] = ["push:agentception", "pull:agentception"]
430 resp = await client.post("/api/identities/agent", json=payload, headers=auth_headers)
431 assert resp.status_code in (200, 201)
432
433
434 # ---------------------------------------------------------------------------
435 # verify_and_authenticate identity_type — unit test
436 # ---------------------------------------------------------------------------
437
438
439 class TestVerifyAuthenticateIdentityType:
440 """Unit test the identity_type param is propagated to the DB row."""
441
442 @pytest.mark.asyncio
443 async def test_identity_type_human_is_default(self) -> None:
444 """When identity_type is omitted it defaults to 'human'."""
445 req = VerifyRequest(
446 challenge_token="a" * 64,
447 public_key_b64="AAEC",
448 signature_b64="AAEC",
449 )
450 assert req.identity_type == "human"
451
452 @pytest.mark.asyncio
453 async def test_identity_type_agent_passes_through(self) -> None:
454 req = VerifyRequest(
455 challenge_token="a" * 64,
456 public_key_b64="AAEC",
457 signature_b64="AAEC",
458 identity_type="agent",
459 )
460 assert req.identity_type == "agent"
File History 1 commit
sha256:9590cee1e0ccd6c76528f005b95d634d80f5019f0dcb7c371e149adc31d1fb65 refactor: enforce gRPC framing on all MWP wire traffic Sonnet 4.6 minor ⚠ 156 days ago