xinmi/scripts/test_operations_etl.py

import copy
import sqlite3
import unittest
from unittest.mock import patch

from weknora_identity_etl import ADMIN_TARGET_ID
from weknora_operations_etl import (KINDS, notification_owner, operations_plan, operations_sql,
                                   reconcile, redact, validate_plan)

USER = "63ce184e-5ab6-5240-8269-14026d501a68"
AGENT = "da4d5f8f-e1c4-574e-b34a-8a324615b603"
KB = "237fa6a0-1e49-58c8-8185-4191d14a9094"
NOW = "2026-09-15T12:00:00"


def fixture():
    rows = {kind: [] for kind in (*KINDS, "campaign", "segment")}
    rows["pointaccount"] = [{"id": 1, "user_id": 1, "balance_points": 110, "updated_at": NOW}]
    rows["pointtransaction"] = [{"id": 1, "user_id": 1, "delta_points": 100, "kind": "credit", "reason": "grant",
                                 "operator_user_id": 1, "metadata": {}, "created_at": NOW},
                                {"id": 2, "user_id": 1, "delta_points": -10, "kind": "debit", "reason": "generation",
                                 "operator_user_id": 1, "metadata": {"usage_event_id": 1}, "created_at": NOW}]
    rows["usageevent"] = [{"id": 1, "user_id": 1, "agent_id": 2, "provider_id": 1, "model": "test",
                           "unit": "token", "quantity": 1, "currency": "USD", "amount_minor": 5000,
                           "llm_call_id": 1, "campaign_id": 1, "run_id": 1, "segment_id": 1,
                           "input_quantity": 1, "output_quantity": 0, "service_type": "llm",
                           "status": "measured", "metadata": {}, "created_at": NOW}]
    rows["campaign"] = [{"id": 1, "owner_user_id": 1}, {"id": 2, "owner_user_id": 8}]
    rows["segment"] = [{"id": 1, "campaign_id": 1}, {"id": 2, "campaign_id": 1}, {"id": 42, "campaign_id": 2}]
    return rows


def plan(rows=None):
    rows = rows or fixture()
    db = sqlite3.connect(":memory:")
    db.execute("PRAGMA query_only=ON")
    try:
        with patch("weknora_operations_etl._read", side_effect=lambda _db, kind: copy.deepcopy(rows[kind])):
            result = operations_plan(db, {"1": ADMIN_TARGET_ID, "8": USER}, {"2": AGENT}, {"3": KB}, {}, {})
    finally:
        db.close()
    result["source_snapshot_sha256"] = "a" * 64
    return result


class OperationsTests(unittest.TestCase):
    def test_mismatch_preserves_authoritative_balance_without_credit(self):
        result = plan()
        self.assertEqual(result["reconciliation"][0]["difference"], 20)
        native = [i for i in result["items"] if i["target_table"] == "workbench_balances"]
        self.assertEqual(native[0]["native_snapshot"]["balance"], 110)
        self.assertEqual(native[0]["history_payload"]["balance_reconciliation"]["difference"], 20)
        self.assertEqual(len([i for i in result["items"] if i["target_table"] == "workbench_transactions"]), 2)
        validate_plan(result)

    def test_missing_account_does_not_create_credit_or_balance(self):
        rows = fixture()
        rows["pointaccount"] = []
        result = plan(rows)
        self.assertEqual(result["reconciliation"][0]["reason"], "missing_source_account")
        self.assertFalse(any(i["target_table"] == "workbench_balances" for i in result["items"]))
        self.assertEqual(sum(i["target_table"] == "workbench_transactions" for i in result["items"]), 2)

    def test_points_use_actual_debits_not_currency_or_token_estimates(self):
        usage = next(i for i in plan()["items"] if i["source_kind"] == "usageevent")
        self.assertEqual(usage["native_snapshot"]["points"], 10)
        self.assertEqual(usage["native_snapshot"]["metadata"]["legacy_amount_minor"], 5000)

    def test_storage_quantity_is_not_tokens(self):
        rows = fixture()
        rows["usageevent"][0].update(unit="byte", input_quantity=999, output_quantity=999)
        usage = next(i["native_snapshot"] for i in plan(rows)["items"] if i["source_kind"] == "usageevent")
        self.assertEqual((usage["input_tokens"], usage["output_tokens"]), (0, 0))

    def test_usage_ownerless_is_admin_archive_only(self):
        rows = fixture()
        rows["usageevent"][0]["user_id"] = 0
        usage = next(i for i in plan(rows)["items"] if i["source_kind"] == "usageevent")
        self.assertEqual(usage["target_table"], "")
        self.assertEqual(usage["owner_id"], "")

    def test_ambiguous_notifications_not_assigned_by_numeric_reference(self):
        ids = {"1": ADMIN_TARGET_ID, "8": USER}
        owner, reason = notification_owner({"type": "failure", "ref_id": 2}, {1: 1, 2: 8}, {2: 1}, ids)
        self.assertEqual(owner, "")
        self.assertTrue(reason)
        owner, reason = notification_owner({"type": "failure", "ref_id": 42}, {1: 1, 2: 8}, {42: 2}, ids)
        self.assertEqual(owner, USER)

    def test_alerts_are_never_broadcast(self):
        owner, _ = notification_owner({"type": "alert", "ref_id": 0}, {}, {}, {"1": ADMIN_TARGET_ID})
        self.assertEqual(owner, "")

    def test_redacted_native_log_preserves_original_in_private_plan(self):
        rows = fixture()
        secret = "sk-123456789abcdefghijklmnop"
        rows["systemlog"] = [{"id": 1, "level": "WARNING", "source": "provider", "message": "api_key=" + secret,
                              "trace_id": "", "status": 500, "created_at": NOW}]
        item = next(i for i in plan(rows)["items"] if i["source_kind"] == "systemlog")
        self.assertIn(secret, item["source_snapshot"]["message"])
        self.assertNotIn(secret, item["native_snapshot"]["message"])
        self.assertEqual(item["native_snapshot"]["level"], "warning")
        self.assertEqual(redact({"api_key": secret})["api_key"], "[REDACTED]")

    def test_native_balance_tampering_rejected(self):
        result = plan()
        result["items"][0]["native_snapshot"]["balance"] += 20
        with self.assertRaisesRegex(ValueError, "balance"):
            validate_plan(result)

    def test_native_owner_tampering_rejected(self):
        result = plan()
        result["items"][0]["native_snapshot"]["user_id"] = USER
        with self.assertRaisesRegex(ValueError, "owner"):
            validate_plan(result)

    def test_duplicate_accounts_rejected(self):
        with self.assertRaisesRegex(ValueError, "Duplicate"):
            reconcile([{"user_id": 1, "balance_points": 0}] * 2, [])

    def test_sql_requires_snapshot_provenance(self):
        result = plan()
        result.pop("source_snapshot_sha256")
        with self.assertRaisesRegex(ValueError, "fingerprint"):
            operations_sql(result)


if __name__ == "__main__":
    unittest.main()