"""Tests for the AWS Bedrock per-profile (``UsedBy`` tag) breakdown on the per-AK spend detail page. Contract under test: * ``_aggregate_by_profile`` reads the ISOLATED ``cmdb_bedrock_profile_facts`` table, groups case-insensitively on ``profile`` (``aiops`` + ``AIOPS`` collapse to one row), folds (profile × service) rows into one entry per profile with a nested ``by_model``, and computes ``percent_of_month`` against the AK's MTD. * ``ingest_bedrock_profile_facts`` is idempotent — re-ingesting the same window replaces rather than duplicates (per-(scope, date) pre-window DELETE), and a shrinking natural-key set evicts stale rows. * The detail endpoint surfaces ``by_profile`` for a Bedrock AK filtered to that AK's owner, and the by-profile sum reconciles to ``mtd_usd``. * The profile facts NEVER touch ``cmdb_cost_facts``, so the existing ``mtd_usd`` / ``by_model`` numbers are unchanged by their presence. """ from __future__ import annotations import uuid from datetime import date, datetime from decimal import Decimal from types import SimpleNamespace from zoneinfo import ZoneInfo import pytest from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.domains.apikey.models import ApiKeyAssignment, ApiKeyProvider from app.domains.apikey.service import ApiKeyService from app.domains.apikey.spend import _aggregate_by_profile from app.domains.cmdb.cost_service import ( ingest_bedrock_profile_facts, ) from app.domains.cmdb.models import ( CmdbBedrockProfileFact, CmdbCostFact, CmdbIamUser, ) from app.identity.person_models import IdentityPerson from app.integrations.models import Integration _BEDROCK_SVC = "Claude Sonnet 4.6 (Amazon Bedrock Edition)" _INTEG = "11111111-1111-1111-1111-111111111111" # ── helpers ──────────────────────────────────────────────────────── async def _seed_profile_fact( db: AsyncSession, *, integration_id: str = _INTEG, scope: str = "841162690310", fact_date: date, owner: str, profile: str, service: str = _BEDROCK_SVC, amount_usd: Decimal, ) -> None: db.add(CmdbBedrockProfileFact( fact_key=uuid.uuid4().hex, date=fact_date, integration_id=integration_id, scope_external_id=scope, owner_email=owner, profile=profile, service=service, amount_usd=amount_usd, amount_native=amount_usd, currency="USD", )) await db.commit() # ── _aggregate_by_profile unit ───────────────────────────────────── @pytest.mark.asyncio async def test_aggregate_by_profile_groups_and_merges_case( db_session: AsyncSession, ): """Three UsedBy values (two of them the same workload under two casings) collapse to two profile rows, sorted by spend desc, each with its model breakdown and the correct share of the AK's MTD.""" today = date(2026, 6, 10) owner = "nizheming@pingcap.com" # investibook-dev across two services + two days. await _seed_profile_fact( db_session, fact_date=today, owner=owner, profile="investibook-dev", amount_usd=Decimal("200.00"), ) await _seed_profile_fact( db_session, fact_date=date(2026, 6, 11), owner=owner, profile="investibook-dev", service="Claude Opus 4.6 (Amazon Bedrock Edition)", amount_usd=Decimal("34.29"), ) # aiops + AIOPS — different casings of one workload, must MERGE. await _seed_profile_fact( db_session, fact_date=today, owner=owner, profile="aiops", amount_usd=Decimal("400.00"), ) await _seed_profile_fact( db_session, fact_date=today, owner=owner, profile="AIOPS", service="Claude Opus 4.6 (Amazon Bedrock Edition)", amount_usd=Decimal("28.97"), ) # A different owner's row — must NOT leak into this owner's split. await _seed_profile_fact( db_session, fact_date=today, owner="someone@pingcap.com", profile="investibook-dev", amount_usd=Decimal("999.00"), ) out = await _aggregate_by_profile( db_session, integration_id=_INTEG, owner_email=owner, month_start=date(2026, 6, 1), month_end=date(2026, 6, 30), month_total_usd=Decimal("663.26"), # 234.29 + 428.97 ) assert [r["profile"] for r in out] == ["aiops", "investibook-dev"] aiops, dev = out[0], out[1] # aiops + AIOPS merged: 400 + 28.97 assert aiops["mtd_usd"] == pytest.approx(428.97) assert dev["mtd_usd"] == pytest.approx(234.29) # Shares of the AK's MTD total. assert aiops["percent_of_month"] == pytest.approx(64.7, abs=0.2) assert dev["percent_of_month"] == pytest.approx(35.3, abs=0.2) # Nested model breakdown: investibook-dev hit two models. dev_models = {m["model"]: m["mtd_usd"] for m in dev["by_model"]} assert dev_models[_BEDROCK_SVC] == pytest.approx(200.00) assert dev_models["Claude Opus 4.6 (Amazon Bedrock Edition)"] == pytest.approx(34.29) @pytest.mark.asyncio async def test_aggregate_by_profile_empty_when_no_rows( db_session: AsyncSession, ): out = await _aggregate_by_profile( db_session, integration_id=_INTEG, owner_email="nobody@pingcap.com", month_start=date(2026, 6, 1), month_end=date(2026, 6, 30), month_total_usd=Decimal("0"), ) assert out == [] # ── ingest idempotency / delete-replace ──────────────────────────── @pytest.mark.asyncio async def test_ingest_bedrock_profile_facts_idempotent( db_session: AsyncSession, ): facts = [ { "date": date(2026, 6, 10), "owner_email": "nizheming@pingcap.com", "profile": "investibook-dev", "service": _BEDROCK_SVC, "amount_usd": 200.0, "amount_native": 200.0, "currency": "USD", "usage_quantity": 1000.0, "usage_unit": "tokens", }, { "date": date(2026, 6, 10), "owner_email": "nizheming@pingcap.com", "profile": "aiops", "service": _BEDROCK_SVC, "amount_usd": 50.0, "amount_native": 50.0, "currency": "USD", "usage_quantity": None, "usage_unit": None, }, ] async def _count() -> int: rows = (await db_session.execute( select(CmdbBedrockProfileFact.id).where( CmdbBedrockProfileFact.integration_id == _INTEG, ) )).all() return len(rows) n1 = await ingest_bedrock_profile_facts( db_session, integration_id=_INTEG, scope_external_id="841162690310", start_date=date(2026, 6, 10), end_date=date(2026, 6, 10), facts=facts, ) assert n1 == 2 assert await _count() == 2 # Re-ingest the SAME window: delete-replace → still 2, no dups. n2 = await ingest_bedrock_profile_facts( db_session, integration_id=_INTEG, scope_external_id="841162690310", start_date=date(2026, 6, 10), end_date=date(2026, 6, 10), facts=facts, ) assert n2 == 2 assert await _count() == 2 # A shrunk natural-key set (aiops retired) evicts the stale row. n3 = await ingest_bedrock_profile_facts( db_session, integration_id=_INTEG, scope_external_id="841162690310", start_date=date(2026, 6, 10), end_date=date(2026, 6, 10), facts=[facts[0]], ) assert n3 == 1 assert await _count() == 1 @pytest.mark.asyncio async def test_ingest_empty_facts_evicts_window(db_session: AsyncSession): await _seed_profile_fact( db_session, fact_date=date(2026, 6, 10), owner="x@pingcap.com", profile="dev", amount_usd=Decimal("5"), ) n = await ingest_bedrock_profile_facts( db_session, integration_id=_INTEG, scope_external_id="841162690310", start_date=date(2026, 6, 10), end_date=date(2026, 6, 10), facts=[], ) assert n == 0 remaining = (await db_session.execute( select(CmdbBedrockProfileFact.id).where( CmdbBedrockProfileFact.integration_id == _INTEG, ) )).all() assert remaining == [] # ── end-to-end through the detail endpoint ───────────────────────── @pytest.mark.asyncio async def test_detail_surfaces_by_profile_for_bedrock_ak( db_session: AsyncSession, ): """A Bedrock AK's detail payload carries ``by_profile`` filtered to that AK's owner, and the by-profile sum equals the model/MTD total — profile facts are additive and never alter ``mtd_usd``.""" provider = ApiKeyProvider( id=uuid.uuid4(), name="AWS BedRock", description="x", is_active=True, sort_order=0, ) integ = Integration( id=uuid.uuid4(), kind="cloud", provider="aws", name="AWS SG", slug=f"aws-{uuid.uuid4().hex[:6]}", enabled=True, is_primary=False, capabilities=[], config_public={}, ) applicant = IdentityPerson( id=uuid.uuid4(), source_system="successfactors", hris_person_id=f"sf-{uuid.uuid4().hex[:8]}", display_name="ni", work_email="nizheming@pingcap.com", employment_status="active", ) db_session.add_all([provider, integ, applicant]) await db_session.commit() iam = CmdbIamUser( id=uuid.uuid4(), integration_id=str(integ.id), user_name="investibook", sub_account_id="", arn="arn:aws:iam::841162690310:user/investibook", tags={"Owner": "nizheming@pingcap.com"}, owner_person_id=str(applicant.id), ) db_session.add(iam) await db_session.commit() assignment = ApiKeyAssignment( id=uuid.uuid4(), provider_id=provider.id, assigned_to=applicant.id, status="active", request_number="AK-PROFTEST", project_name="investibook", reason="t", monthly_budget_usd=Decimal("1500.00"), iam_user_arn=iam.arn, iam_integration_id=str(integ.id), iam_user_id=iam.id, ) db_session.add(assignment) await db_session.commit() today = datetime.now(ZoneInfo("Asia/Shanghai")).date() month_start = today.replace(day=1) # cmdb_cost_facts: the AK's owner total ($250) — drives mtd_usd. db_session.add(CmdbCostFact( fact_key=uuid.uuid4().hex, date=today, integration_id=str(integ.id), service=_BEDROCK_SVC, amount_usd=Decimal("250.00"), amount_native=Decimal("250.00"), currency="USD", resource_id=( f"aws-principal:Owner=nizheming@pingcap.com|service={_BEDROCK_SVC}" ), tags={"Owner": "nizheming@pingcap.com"}, auto_owner_source="scope_owner", )) await db_session.commit() # Isolated profile facts split the same $250: dev 150 + aiops 100. await _seed_profile_fact( db_session, integration_id=str(integ.id), fact_date=today, owner="nizheming@pingcap.com", profile="investibook-dev", amount_usd=Decimal("150.00"), ) await _seed_profile_fact( db_session, integration_id=str(integ.id), fact_date=today, owner="nizheming@pingcap.com", profile="aiops", amount_usd=Decimal("100.00"), ) out = await ApiKeyService(db_session).get_assignment_spend_detail( assignment.id, current_user=SimpleNamespace(id=applicant.id, role_names=["employee"]), ) assert out["bucket"] == "bedrock" assert out["mtd_usd"] == pytest.approx(250.00) by_profile = out["by_profile"] assert [r["profile"] for r in by_profile] == ["investibook-dev", "aiops"] assert sum(r["mtd_usd"] for r in by_profile) == pytest.approx(250.00) # by_model (cmdb_cost_facts) is unaffected by the profile facts. assert sum(r["mtd_usd"] for r in out["by_model"]) == pytest.approx(250.00) _ = month_start def test_driver_profile_tag_key_is_verbatim_not_iamprincipal_prefixed(): """Regression: the profile tag (`UsedBy`) is a PLAIN resource cost-allocation tag — CE key `UsedBy` — NOT an IAM-principal tag. The owner tag (`Owner`) IS a principal tag → `iamPrincipal/Owner`. The driver must prefix the owner but use the profile key verbatim; prefixing the profile as `iamPrincipal/UsedBy` returns empty for every row and collapses the breakdown to `(unallocated)` (the 2026-06-17 prod backfill bug). This drives the real CE GroupBy construction with a fake boto client (the other tests mock the whole driver and so never exercised the tag-key namespacing). """ from app.integrations.driver_base import CloudScopeRef from app.integrations.drivers.aws_cloud import AwsCloudDriver recorded: list[list[dict]] = [] class _FakeCE: def get_cost_and_usage(self, **kw): gb = kw.get("GroupBy", []) recorded.append(gb) tag_keys = [g["Key"] for g in gb if g.get("Type") == "TAG"] if tag_keys == ["iamPrincipal/Owner"]: # owner discovery return {"ResultsByTime": [{"TimePeriod": {"Start": "2026-06-10"}, "Groups": [{"Keys": ["iamPrincipal/Owner$nizheming@pingcap.com"], "Metrics": {"UnblendedCost": {"Amount": "10", "Unit": "USD"}}}]}]} # per-owner profile × service breakdown return {"ResultsByTime": [{"TimePeriod": {"Start": "2026-06-10"}, "Groups": [{"Keys": ["UsedBy$investibook-dev", "Claude Sonnet 4.6 (Amazon Bedrock Edition)"], "Metrics": {"UnblendedCost": {"Amount": "5", "Unit": "USD"}, "UsageQuantity": {"Amount": "1", "Unit": "tokens"}}}]}]} d = AwsCloudDriver.__new__(AwsCloudDriver) d.integration_id = "x" d._boto_client = lambda *a, **k: _FakeCE() # type: ignore[method-assign] scope = CloudScopeRef( scope_type="aws_account", external_id="841162690310", display_name="x", ) facts = d._list_bedrock_profile_facts_sync( date(2026, 6, 1), date(2026, 6, 16), scope, "Owner", "UsedBy", None, ) all_tag_keys = {g["Key"] for gb in recorded for g in gb if g.get("Type") == "TAG"} assert "UsedBy" in all_tag_keys, "profile tag must be queried verbatim" assert "iamPrincipal/UsedBy" not in all_tag_keys, "must NOT prefix profile tag" assert "iamPrincipal/Owner" in all_tag_keys, "owner tag IS principal-prefixed" # The profile value resolved (not collapsed to (unallocated)). assert facts and facts[0]["profile"] == "investibook-dev" assert facts[0]["owner_email"] == "nizheming@pingcap.com" @pytest.mark.asyncio async def test_sync_filters_non_bedrock_services( db_session: AsyncSession, db_engine, monkeypatch, ): """The owner cost-allocation tag is an IAM-principal tag, so the CE fan-out returns ALL the owner's services. The sync MUST drop the non-Bedrock rows (per the AI whitelist) before storing, so the per-profile sum reconciles to the Bedrock-only ``mtd_usd`` instead of inflating it with the principal's EC2/S3 spend. Regression guard for the review's HIGH finding. """ from app.core import database as _db from app.domains.cmdb import cost_service as _cs integ = Integration( id=uuid.uuid4(), kind="cloud", provider="aws", name="AWS X", slug="aws-x", enabled=True, is_primary=False, capabilities=[], config_public={}, ) db_session.add(integ) await db_session.commit() # integration row must exist before scope FK from app.integrations.models import PlatformSetting db_session.add_all([ PlatformSetting(key="cost_facts.aws_principal_tag.aws-x", value="Owner"), PlatformSetting(key="cost_facts.aws_profile_tag.aws-x", value="UsedBy"), ]) from app.domains.cmdb.models import CmdbCloudScope db_session.add(CmdbCloudScope( id=uuid.uuid4(), integration_id=str(integ.id), scope_type="aws_account", external_id="841162690310", display_name="AWS X", status="active", )) await db_session.commit() # Fake driver: returns one EC2 (non-Bedrock) row that MUST be dropped. mixed = [ {"date": date(2026, 6, 10), "owner_email": "n@pingcap.com", "profile": "investibook-dev", "service": _BEDROCK_SVC, "amount_usd": 150.0, "amount_native": 150.0, "currency": "USD", "usage_quantity": None, "usage_unit": None}, {"date": date(2026, 6, 10), "owner_email": "n@pingcap.com", "profile": "investibook-dev", "service": "Amazon EC2", "amount_usd": 999.0, "amount_native": 999.0, "currency": "USD", "usage_quantity": None, "usage_unit": None}, {"date": date(2026, 6, 10), "owner_email": "n@pingcap.com", "profile": "aiops", "service": "Amazon Bedrock", "amount_usd": 100.0, "amount_native": 100.0, "currency": "USD", "usage_quantity": None, "usage_unit": None}, ] class _FakeDriver: requires_cny_fx = False async def list_bedrock_profile_facts( self, start, end, *, scope=None, owner_tag_key, profile_tag_key, cny_per_usd=None, ): return list(mixed) async def close(self): return None async def _fake_load(session, integration): return _FakeDriver() async def _fake_bq(session, integration): return None # Bind the module-level session factory to the test engine (the # sync opens its own session), and stub the driver + BQ resolution. from sqlalchemy.ext.asyncio import async_sessionmaker monkeypatch.setattr( _db, "async_session_factory", async_sessionmaker(db_engine, expire_on_commit=False), ) monkeypatch.setattr( "app.domains.cmdb.lifecycle_service._load_cloud_driver", _fake_load, ) monkeypatch.setattr( "app.domains.cmdb.lifecycle_service._resolve_aws_bq_billing", _fake_bq, ) res = await _cs.sync_bedrock_profile_facts_for_integration( str(integ.id), start_date=date(2026, 6, 10), end_date=date(2026, 6, 10), ) assert res["status"] == "ok" # Only the two Bedrock rows landed — the $999 EC2 row was filtered. rows = (await db_session.execute( select( CmdbBedrockProfileFact.service, CmdbBedrockProfileFact.amount_usd, ).where(CmdbBedrockProfileFact.integration_id == str(integ.id)) )).all() services = {r[0] for r in rows} assert "Amazon EC2" not in services assert services == {_BEDROCK_SVC, "Amazon Bedrock"} assert sum(float(r[1]) for r in rows) == pytest.approx(250.0)