diff --git a/src/lightspeed_agent/dcr/service.py b/src/lightspeed_agent/dcr/service.py index a28f011b..913b9068 100644 --- a/src/lightspeed_agent/dcr/service.py +++ b/src/lightspeed_agent/dcr/service.py @@ -172,6 +172,11 @@ async def register_client( error_description=f"Invalid Order ID: {claims.order_id}", ) + # Step 3b: Backfill entitlement account_id if empty + await self._procurement_service.backfill_entitlement_account_id( + claims.order_id, claims.account_id + ) + # Step 4: Check if client already exists for this order existing_client = await self._client_repository.get_by_order_id(claims.order_id) if existing_client: diff --git a/src/lightspeed_agent/marketplace/service.py b/src/lightspeed_agent/marketplace/service.py index 680c927e..c2fafeae 100644 --- a/src/lightspeed_agent/marketplace/service.py +++ b/src/lightspeed_agent/marketplace/service.py @@ -459,6 +459,37 @@ async def _handle_offer_ended(self, event: ProcurementEvent) -> None: if event.entitlement: logger.info("Offer ended: %s", event.entitlement.id) + # Public helpers + + async def backfill_entitlement_account_id(self, order_id: str, account_id: str) -> None: + """Backfill an entitlement's account_id if it is currently empty. + + Called during DCR registration when the JWT provides an authoritative + account_id that may be missing from the entitlement record (Pub/Sub + events often arrive without account info). + + Args: + order_id: The entitlement/order ID. + account_id: The account ID to set (from the DCR JWT ``sub`` claim). + """ + if not account_id: + return + + entitlement = await self._entitlement_repo.get(order_id) + if not entitlement: + return + + if entitlement.account_id: + return + + entitlement.account_id = account_id + await self._entitlement_repo.update(entitlement) + logger.info( + "Backfilled account_id %s on entitlement %s from DCR", + account_id, + order_id, + ) + # Procurement API operations async def _resolve_account_id(self, entitlement_id: str, event: ProcurementEvent) -> str | None: @@ -519,6 +550,11 @@ async def _resolve_account_id(self, entitlement_id: str, event: ProcurementEvent "Entitlement %s not found in Procurement API (404)", entitlement_id, ) + elif response.status_code >= 500: + raise RuntimeError( + f"Server error resolving account for entitlement {entitlement_id}: " + f"HTTP {response.status_code} — {response.text}" + ) else: logger.warning( "Failed to fetch entitlement %s: HTTP %s — %s", diff --git a/tests/test_dcr.py b/tests/test_dcr.py index 21ad31f8..3e874c0f 100644 --- a/tests/test_dcr.py +++ b/tests/test_dcr.py @@ -172,6 +172,74 @@ async def test_get_client(self, service): assert client.account_id == "valid-account-123" +class TestDCRServiceBackfill: + """Tests for DCR service backfilling entitlement account_id.""" + + @pytest_asyncio.fixture + async def service_and_repo(self, db_session): + """Create a DCR service with an entitlement that has empty account_id.""" + entitlement_repo = EntitlementRepository() + client_repo = DCRClientRepository() + procurement_service = ProcurementService( + entitlement_repo=entitlement_repo, + ) + + entitlement = Entitlement( + id="backfill-order-789", + account_id="", + provider_id="provider-456", + state=EntitlementState.ACTIVE, + ) + await entitlement_repo.create(entitlement) + + svc = DCRService( + procurement_service=procurement_service, + client_repository=client_repo, + ) + return svc, entitlement_repo + + @pytest.mark.asyncio + async def test_register_client_backfills_entitlement_account_id(self, service_and_repo): + """Test that DCR registration backfills empty entitlement account_id.""" + from lightspeed_agent.dcr.gma_client import GMAClientResponse + + service, entitlement_repo = service_and_repo + + claims = GoogleJWTClaims( + iss="https://accounts.google.com", + iat=int(time.time()), + exp=int(time.time()) + 3600, + aud="https://example.com", + sub="backfill-account-123", + google={"order": "backfill-order-789"}, + auth_app_redirect_uris=["https://example.com/callback"], + ) + + mock_gma = AsyncMock() + mock_gma.create_tenant.return_value = GMAClientResponse( + client_id="new-client-id", + client_secret="new-secret", + name="lightspeed-backfill-order-789", + ) + service._gma_client = mock_gma + + # Mock JWT validation to return our claims + with patch.object( + service._jwt_validator, + "validate_software_statement", + new_callable=AsyncMock, + return_value=claims, + ): + result = await service.register_client( + DCRRequest(software_statement="valid-jwt") + ) + + assert isinstance(result, DCRResponse) + + ent = await entitlement_repo.get("backfill-order-789") + assert ent.account_id == "backfill-account-123" + + class TestDCRServiceEncryptionValidation: """Tests for DCR service encryption key validation.""" diff --git a/tests/test_marketplace.py b/tests/test_marketplace.py index 09c637b0..229ad623 100644 --- a/tests/test_marketplace.py +++ b/tests/test_marketplace.py @@ -373,8 +373,72 @@ async def test_entitlement_creation_idempotent(self, service): mock_settings.skip_pubsub_oidc_verification = False await service.process_event(event) + # --- backfill_entitlement_account_id tests --- + + @pytest.mark.asyncio + async def test_backfill_entitlement_account_id(self, service): + """Test backfill updates empty account_id.""" + ent = Entitlement( + id="order-backfill", + account_id="", + state=EntitlementState.ACTIVE, + provider_id="provider-1", + ) + await service._entitlement_repo.create(ent) + + await service.backfill_entitlement_account_id("order-backfill", "account-filled") + + updated = await service._entitlement_repo.get("order-backfill") + assert updated.account_id == "account-filled" + + @pytest.mark.asyncio + async def test_backfill_entitlement_account_id_no_overwrite(self, service): + """Test backfill does not overwrite existing account_id.""" + ent = Entitlement( + id="order-existing-acct", + account_id="original-account", + state=EntitlementState.ACTIVE, + provider_id="provider-1", + ) + await service._entitlement_repo.create(ent) + + await service.backfill_entitlement_account_id("order-existing-acct", "new-account") + + updated = await service._entitlement_repo.get("order-existing-acct") + assert updated.account_id == "original-account" + + @pytest.mark.asyncio + async def test_backfill_entitlement_account_id_not_found(self, service): + """Test backfill on non-existent order does not raise.""" + await service.backfill_entitlement_account_id("order-nonexistent", "account-1") + # --- _resolve_account_id tests --- + @pytest.mark.asyncio + async def test_resolve_account_id_raises_on_5xx(self, service): + """Test _resolve_account_id raises RuntimeError on 5xx server errors.""" + event = ProcurementEvent( + event_id="event-1", + event_type=ProcurementEventType.ENTITLEMENT_CREATION_REQUESTED, + provider_id="provider-1", + entitlement=EntitlementInfo(id="order-1"), + ) + + mock_response = httpx.Response( + status_code=500, + text="Internal Server Error", + request=httpx.Request("GET", "https://example.com"), + ) + with ( + patch.object(service, "_settings") as mock_settings, + patch.object(service, "_get_auth_headers", return_value={}), + patch("httpx.AsyncClient.get", new_callable=AsyncMock, return_value=mock_response), + ): + mock_settings.google_cloud_project = "test-project" + mock_settings.skip_pubsub_oidc_verification = False + with pytest.raises(RuntimeError, match="Server error resolving account"): + await service._resolve_account_id("order-1", event) + @pytest.mark.asyncio async def test_resolve_account_id_from_event(self, service): """Test _resolve_account_id returns account ID from event payload."""