adcloud-server/tests/test_projection.py

import copy
from dataclasses import replace

import pytest
from sqlalchemy import select
from sqlalchemy.dialects import postgresql
from sqlalchemy.schema import CreateTable

from adcloud_server.database import Database, outbox, profiles
from adcloud_server.identity_models import ProfilePatch, StitchRequest
from adcloud_server.identity_service import IdentityService
from adcloud_server.projection import ProfileProjector
from adcloud_server.settings import Settings


class RecoverableStore:
    def __init__(self):
        self.snapshots = {}
        self.unavailable = False

    def put(self, tenant, master, revision, document):
        if self.unavailable:
            raise OSError("Simulated storage outage")
        self.snapshots[tenant, master, revision] = copy.deepcopy(document)

    def get(self, tenant, master, revision):
        if self.unavailable:
            raise OSError("Simulated storage outage")
        return self.snapshots.get((tenant, master, revision))


@pytest.fixture
def components(tmp_path):
    settings = Settings(
        database_url="sqlite:///" + str(tmp_path / "profiles.db"),
        identity_secret="x" * 48,
        profile_backend="cassandra",
    )
    db = Database(settings.database_url, settings.tenant_id)
    db.initialize()
    store = RecoverableStore()
    service = IdentityService(settings, db, store)
    try:
        yield service, ProfileProjector(db, store), store
    finally:
        db.close()


def test_wide_table_outage_preserves_immediate_stitched_profile_and_retries(components):
    service, worker, store = components
    anon = service.anonymous()
    result = service.stitch(
        anon["anonymous_token"],
        StitchRequest(
            identifiers=[{"kind": "email", "value": "alice@example.com"}],
            attributes={"interest": "running"},
            tags=["runner"],
        ),
    )
    assert service.profile(result["master_id"])["projection_status"] == "pending"
    store.unavailable = True
    assert service.profile(result["master_id"])["attributes"] == {"interest": "running"}
    with pytest.raises(OSError):
        worker.flush()
    with service.db.transaction() as conn:
        pending = (
            conn.execute(select(outbox).where(outbox.c.delivered == 0)).mappings().all()
        )
        assert any(job["revision"] == result["revision"] for job in pending)
    store.unavailable = False
    worker.flush()
    current = service.profile(result["master_id"])
    assert current["projection_status"] == "current"
    assert current["tags"] == ["runner"]
    assert worker.flush() == 0


def test_out_of_order_projection_cannot_resurrect_deleted_attributes(components):
    service, worker, store = components
    anon = service.anonymous()
    first = service.patch_profile(
        anon["master_id"], ProfilePatch(attributes={"remove": "old", "keep": 1})
    )
    store.put(
        service.tenant,
        anon["master_id"],
        first["revision"],
        {"attributes": first["attributes"], "tags": []},
    )
    last = service.patch_profile(
        anon["master_id"], ProfilePatch(attributes={"remove": None})
    )
    # A stale snapshot never satisfies the latest revision lookup.
    assert service.profile(anon["master_id"])["attributes"] == {"keep": 1}
    assert service.profile(anon["master_id"])["projection_status"] == "pending"
    worker.flush()
    assert service.profile(anon["master_id"])["revision"] == last["revision"]
    assert service.profile(anon["master_id"])["attributes"] == {"keep": 1}


def test_relational_to_wide_table_switch_backfills_existing_profiles(components):
    service, worker, _store = components
    relational = IdentityService(
        replace(service.settings, profile_backend="relational"), service.db
    )
    anon = relational.anonymous()
    relational.patch_profile(
        anon["master_id"], ProfilePatch(attributes={"nested": {"count": 3}})
    )
    with service.db.transaction() as conn:
        assert conn.execute(select(outbox)).first() is None
    assert worker.backfill() == 1
    worker.flush()
    assert service.profile(anon["master_id"])["projection_status"] == "current"
    assert service.profile(anon["master_id"])["attributes"] == {"nested": {"count": 3}}


def test_merge_failure_never_leaves_partial_outbox_records(components, monkeypatch):
    service, worker, _store = components
    a = service.anonymous()
    service.stitch(
        a["anonymous_token"],
        StitchRequest(identifiers=[{"kind": "email", "value": "alice@example.com"}]),
    )
    b = service.anonymous(rotate=True)
    worker.flush()

    def fail(*args):
        raise OSError("Simulated transaction failure")

    monkeypatch.setattr(service, "_queue", fail)
    with pytest.raises(OSError):
        service.stitch(
            b["anonymous_token"],
            StitchRequest(
                identifiers=[{"kind": "email", "value": "alice@example.com"}]
            ),
        )
    assert service.profile(b["master_id"])["master_id"] == b["master_id"]
    assert worker.flush() == 0
    assert len(service.graph(a["master_id"])["member_ids"]) == 1


def test_postgres_document_column_is_jsonb():
    ddl = str(CreateTable(profiles).compile(dialect=postgresql.dialect()))
    assert "document JSONB NOT NULL" in ddl