xinmi/scripts/test_weknora_model_etl.py

"""Safe planner, source routing, and generated-SQL model migration contracts."""
import base64
import copy
import json
import sqlite3
import sys
import unittest
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
from weknora_model_etl import (LEDGER, TARGET_TENANT, model_id, model_plan, model_sql,
                               model_verification_sql, read_credentials, resolve_agent_model)


def fake_encryption(value):
    # This verifies the encoder boundary only; it is never a production cipher.
    return "enc:v1:" + base64.urlsafe_b64encode(b"x" * (12 + len(value.encode()) + 16)).decode().rstrip("=")


class ModelTests(unittest.TestCase):
    def setUp(self):
        self.db = sqlite3.connect(":memory:")
        self.addCleanup(self.db.close)
        self.db.executescript("""
            CREATE TABLE modelprovider (id INTEGER,name TEXT,base_url TEXT,api_key TEXT,models TEXT,
                default_model TEXT,builtin INTEGER,enabled INTEGER,organization_id INTEGER);
            CREATE TABLE modelcatalog (id INTEGER,provider_id INTEGER,model_name TEXT,display_name TEXT,
                capabilities TEXT,capability_status TEXT,enabled INTEGER,builtin INTEGER,rate_version TEXT,
                input_rate REAL,output_rate REAL);
            CREATE TABLE appsetting (key TEXT,value TEXT);
            INSERT INTO modelprovider VALUES (3,'deepseek','https://api.deepseek.com/v1','source-secret-UNIQUE',
                '["deepseek-flash","deepseek-v4-pro"]','deepseek-flash',0,1,1);
            INSERT INTO modelprovider VALUES (4,'images','https://images.invalid/v1','image-secret-UNIQUE',
                '["gpt-image-2"]','gpt-image-2',0,1,1);
            INSERT INTO modelcatalog VALUES (4,3,'deepseek-v4-flash','Disabled','[]','untested',0,0,'r1',1,2);
            INSERT INTO modelcatalog VALUES (6,3,'deepseek-v4-pro','Pro','[]','untested',1,0,'r1',1,2);
            INSERT INTO modelcatalog VALUES (7,3,'deepseek-flash','Flash','[]','untested',1,0,'r1',1,2);
            INSERT INTO modelcatalog VALUES (8,4,'gpt-image-2','Image','[]','untested',1,0,'r1',1,2);
            INSERT INTO appsetting VALUES ('default_model_provider_id','3');
        """)

    def plan(self):
        plan = model_plan(self.db)
        plan["source_snapshot_sha256"] = "a" * 64
        return plan

    def test_reviewed_models_only_without_native_image_generation_guess(self):
        plan = self.plan()
        self.assertEqual(len(plan["models"]), 2)
        self.assertEqual(plan["catalog_map"], {"6": model_id(3, "deepseek-v4-pro"), "7": model_id(3, "deepseek-flash")})
        self.assertEqual(plan["default_gen_model_id"], model_id(3, "deepseek-flash"))
        self.assertEqual({item["reason"] for item in plan["skipped"]},
                         {"source_configuration_disabled", "image_generation_has_no_native_model_type"})
        for item in plan["models"]:
            self.assertFalse(item["is_default"])
            self.assertFalse(item["is_builtin"])
            self.assertEqual(item["source_catalog"]["capability_status"], "untested")
            self.assertIn("no_paid_probe_performed", item["compatibility"])
        self.assertNotIn("secret-UNIQUE", json.dumps(plan))

    def test_planner_reads_no_credentials_and_never_writes(self):
        def authorize(action, first, second, database, trigger):
            if action == sqlite3.SQLITE_READ and first == "modelprovider" and second == "api_key":
                return sqlite3.SQLITE_DENY
            return sqlite3.SQLITE_OK
        self.db.set_authorizer(authorize)
        self.db.execute("PRAGMA query_only=ON")
        self.assertEqual(len(self.plan()["models"]), 2)

    def test_source_runtime_fallback_is_explicit_not_catalog_guess(self):
        result = resolve_agent_model(self.plan(), {"model_provider_id": None, "model_name": "deepseek-chat"})
        self.assertEqual(result["effective_provider_id"], 3)
        self.assertEqual(result["effective_model_name"], "deepseek-flash")
        self.assertEqual(result["requested_model_name"], "deepseek-chat")
        self.assertTrue(result["provider_fallback"])
        self.assertTrue(result["model_fallback"])
        direct = resolve_agent_model(self.plan(), {"model_provider_id": 3, "model_name": "deepseek-v4-pro"})
        self.assertEqual(direct["target_model_id"], model_id(3, "deepseek-v4-pro"))
        self.assertFalse(direct["model_fallback"])

    def test_unsupported_effective_agent_model_does_not_silently_fall_back(self):
        with self.assertRaisesRegex(ValueError, "no reviewed native mapping"):
            resolve_agent_model(self.plan(), {"model_provider_id": 4, "model_name": "gpt-image-2"})

    def test_global_default_uses_source_builtin_and_id_order_when_setting_absent(self):
        self.db.execute("DELETE FROM appsetting")
        self.assertEqual(self.plan()["default_provider_id"], 3)
        self.db.execute("UPDATE modelprovider SET builtin=1 WHERE id=4")
        plan = self.plan()
        self.assertEqual(plan["default_provider_id"], 4)
        self.assertIsNone(plan["default_gen_model_id"])
        with self.assertRaisesRegex(ValueError, "no reviewed native mapping"):
            resolve_agent_model(plan, {"model_provider_id": None})

    def test_provider_default_and_first_configured_model_match_legacy_runtime(self):
        self.db.execute("UPDATE modelprovider SET default_model='' WHERE id=3")
        resolved = resolve_agent_model(self.plan(), {"model_provider_id": 3, "model_name": ""})
        self.assertEqual(resolved["effective_model_name"], "deepseek-flash")

    def test_credentials_read_only_for_imported_models(self):
        def authorize(action, first, second, database, trigger):
            return sqlite3.SQLITE_OK
        self.db.set_authorizer(authorize)
        credentials = read_credentials(self.db, self.plan())
        self.assertEqual(set(credentials), {3})
        self.assertEqual(credentials[3], "source-secret-UNIQUE")

    def test_source_keychain_reference_needs_resolver(self):
        self.db.execute("UPDATE modelprovider SET api_key='keychain:approved:provider-3' WHERE id=3")
        with self.assertRaisesRegex(ValueError, "resolver"):
            read_credentials(self.db, self.plan())
        credentials = read_credentials(self.db, self.plan(), lambda value, provider_id: "resolved-in-memory")
        self.assertEqual(credentials, {3: "resolved-in-memory"})

    def test_no_plaintext_sql_or_credential_ledger_and_no_existing_model_updates(self):
        plan = self.plan()
        sql = model_sql(plan, read_credentials(self.db, plan), fake_encryption)
        self.assertNotIn("source-secret-UNIQUE", sql)
        self.assertNotIn("image-secret-UNIQUE", sql)
        self.assertNotIn("UPDATE models", sql)
        self.assertNotIn("ON CONFLICT", sql)
        self.assertIn("item - 'stored_parameters'", sql)
        self.assertIn("current_database() <> 'yawan_native_acceptance_20260915'", sql)
        self.assertIn("Model ID collision without migration provenance", sql)
        self.assertIn("existing_model.deleted_at IS NOT NULL", sql)
        self.assertIn("mapped.source_snapshot IS DISTINCT FROM", sql)
        self.assertIn("SET LOCAL log_min_error_statement = 'panic'", sql)
        self.assertIn(LEDGER, sql)
        self.assertNotIn("stored_parameters", json.dumps(plan))

    def test_native_cipher_required_and_callback_errors_redacted(self):
        plan = self.plan()
        for encrypt in (lambda x: x, lambda x: "enc:v1:invalid", lambda x: "enc:v1:has padding=="):
            with self.subTest(encrypt=encrypt), self.assertRaisesRegex(ValueError, "encryption failed"):
                model_sql(plan, {3: "source-secret-UNIQUE"}, encrypt)
        def broken(value):
            raise RuntimeError(value)
        try:
            model_sql(plan, {3: "source-secret-UNIQUE"}, broken)
        except ValueError as error:
            self.assertNotIn("source-secret", str(error))
            self.assertTrue(error.__suppress_context__)
        else:
            self.fail("encryption failure accepted")

    def test_missing_secrets_and_snapshot_are_rejected(self):
        plan = self.plan()
        for secret in (None, "", "keychain:reference", "enc:v1:reference", "env:reference"):
            with self.subTest(secret=secret), self.assertRaises(ValueError):
                model_sql(plan, {3: secret}, fake_encryption)
        del plan["source_snapshot_sha256"]
        with self.assertRaisesRegex(ValueError, "fingerprint"):
            model_sql(plan, {3: "in-memory"}, fake_encryption)

    def test_tenant_protocol_and_organization_guards(self):
        for tenant in (1, str(TARGET_TENANT), True):
            with self.subTest(tenant=tenant), self.assertRaises(ValueError):
                model_plan(self.db, tenant)
        self.db.execute("UPDATE modelprovider SET organization_id=2 WHERE id=4")
        with self.assertRaisesRegex(ValueError, "organization"):
            self.plan()

    def test_unreviewed_endpoint_and_model_explicitly_skipped(self):
        self.db.execute("UPDATE modelprovider SET base_url='https://other.invalid/v1' WHERE id=3")
        plan = self.plan()
        self.assertEqual(plan["models"], [])
        self.assertIn("model_protocol_not_reviewed", {s["reason"] for s in plan["skipped"]})

    def test_private_url_never_enters_plan(self):
        self.db.execute("UPDATE modelprovider SET base_url='https://example.invalid/v1?key=private' WHERE id=3")
        with self.assertRaisesRegex(ValueError, "private components"):
            self.plan()

    def test_ids_and_input_mapping_are_deterministic(self):
        self.assertEqual(self.plan(), self.plan())
        self.assertNotEqual(model_id(3, "deepseek-flash"), model_id(4, "deepseek-flash"))
        self.assertNotEqual(model_id(3, "deepseek-flash"), model_id(3, "deepseek-flash", 10000))
        plan = copy.deepcopy(self.plan())
        plan["provider_model_map"]["3:deepseek-flash"] = "wrong"
        with self.assertRaisesRegex(ValueError, "mapping"):
            resolve_agent_model(plan, {})

    def test_sql_rejects_scope_and_native_configuration_tampering(self):
        mutations = [lambda p: p.update(target_tenant_id=10000),
                     lambda p: p.update(target_database="production"),
                     lambda p: p["models"][0].update(is_default=True),
                     lambda p: p["models"][0].update(is_builtin=True),
                     lambda p: p["models"][0]["parameters"].update(api_key="must-not-be-in-plan")]
        for mutation in mutations:
            plan = copy.deepcopy(self.plan())
            mutation(plan)
            with self.subTest(mutation=mutation), self.assertRaises(ValueError):
                model_sql(plan, {3: "in-memory"}, fake_encryption)

    def test_verification_emits_only_secret_presence(self):
        sql = model_verification_sql()
        self.assertIn("LIKE 'enc:v1:%'", sql)
        self.assertNotIn("row_to_json", sql)


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