tests/test_model_loader.py

import tempfile
import unittest
import zipfile
from pathlib import Path
from unittest.mock import patch

from bucklake_ai.config import Settings
from bucklake_ai.errors import ModelArtifactNotFound
from bucklake_ai.model_loader import ModelLoader
from bucklake_ai.model_build import get_model_custom_objects


class StubTensor:
    def __init__(self, name, shape):
        self.name = name
        self.shape = shape


class StubModel:
    output_names = [
        "cum_close_log_returns",
        "open_log_gaps",
        "upper_spans",
        "lower_spans",
        "rendered_ohlc_prices",
        "stock_direction",
        "risk_score",
        "text_gate",
    ]

    inputs = [
        StubTensor("text_input", (None, 768)),
        StubTensor("pos_input", (None, 1024)),
        StubTensor("entity_input", (None, 1024)),
        StubTensor("sentiment_input", (None, 1)),
        StubTensor("index_us_stock_index_INX", (None, 30, 6)),
        StubTensor("index_us_stock_index_DJ", (None, 30, 6)),
        StubTensor("index_us_stock_index_IXIC", (None, 30, 6)),
        StubTensor("index_us_stock_index_NDX", (None, 30, 6)),
        StubTensor("stock_input", (None, 30, 6)),
        StubTensor("pool_symbol_input", (None, 768)),
        StubTensor("pool_sector_input", (None, 768)),
        StubTensor("pool_industry_input", (None, 768)),
        StubTensor("pool_market_input", (None, 768)),
        StubTensor("pool_stats_input", (None, 8)),
    ]


class StubOutput:
    def __init__(self, shape):
        self.shape = shape


class StubModelV2(StubModel):
    output_names = StubModel.output_names + [
        "index_inx_path_output",
        "index_dj_path_output",
        "index_ixic_path_output",
        "index_ndx_path_output",
    ]
    outputs = [StubOutput((None, 3, 1)) for _ in range(4)] + [
        StubOutput((None, 3, 4)),
        StubOutput((None, 3, 1)),
        StubOutput((None, 3, 1)),
        StubOutput((None, 3, 1)),
        StubOutput((None, 3, 4)),
        StubOutput((None, 3, 4)),
        StubOutput((None, 3, 4)),
        StubOutput((None, 3, 4)),
    ]


class ModelLoaderTests(unittest.TestCase):
    def test_custom_objects_include_stack_features_for_round41_artifact(self):
        self.assertIn("StackFeatures", get_model_custom_objects())

    def test_anchored_path_prefers_existing_local_model_path(self):
        with tempfile.TemporaryDirectory() as tmp:
            root_dir = Path(tmp)
            cache_dir = root_dir / ".cache" / "huggingface"
            local_path = root_dir / "local_anchored_path_v1.keras"
            local_path.write_text("local", encoding="utf-8")
            settings = Settings(
                app_name="BuckLakeAI",
                app_version="2.0.0",
                model_version="anchored-path-v1",
                input_contract_version="v2.1.0",
                root_dir=root_dir,
                hf_repo_id="parkerjj/BuckLake-Stock-Model",
                hf_model_filename="anchored_path_v1.keras",
                hf_scaler_filename="scalers_v21_rolling_7d.json",
                hf_cache_dir=cache_dir,
                hf_token=None,
                model_weights_path=local_path,
                text_encoder_model="test-encoder",
                text_encoder_device="cpu",
                text_encoder_max_seq_length=512,
                enable_finbert_sentiment=False,
                finbert_model_name="ProsusAI/finbert",
                preload_model=False,
                preload_text_encoder=False,
                scaler_artifact_path=root_dir / "scalers.json",
            )

            self.assertEqual(ModelLoader(settings)._resolve_model_path(), local_path)

    def test_anchored_path_explicit_local_model_path_wins_over_cache(self):
        with tempfile.TemporaryDirectory() as tmp:
            root_dir = Path(tmp)
            cache_dir = root_dir / ".cache" / "huggingface"
            cache_dir.mkdir(parents=True)
            cached_path = cache_dir / "anchored_path_v1.keras"
            local_path = root_dir / "local_anchored_path_v1.keras"
            cached_path.write_text("cached", encoding="utf-8")
            local_path.write_text("local", encoding="utf-8")
            settings = Settings(
                app_name="BuckLakeAI",
                app_version="2.0.0",
                model_version="anchored-path-v1",
                input_contract_version="v2.1.0",
                root_dir=root_dir,
                hf_repo_id="parkerjj/BuckLake-Stock-Model",
                hf_model_filename="anchored_path_v1.keras",
                hf_scaler_filename="scalers_v21_rolling_7d.json",
                hf_cache_dir=cache_dir,
                hf_token=None,
                model_weights_path=local_path,
                text_encoder_model="test-encoder",
                text_encoder_device="cpu",
                text_encoder_max_seq_length=512,
                enable_finbert_sentiment=False,
                finbert_model_name="ProsusAI/finbert",
                preload_model=False,
                preload_text_encoder=False,
                scaler_artifact_path=root_dir / "scalers.json",
            )

            self.assertEqual(ModelLoader(settings)._resolve_model_path(), local_path)

    def test_anchored_path_downloads_round7_model_filename(self):
        with tempfile.TemporaryDirectory() as tmp:
            root_dir = Path(tmp)
            cache_dir = root_dir / ".cache" / "huggingface"
            downloaded_path = root_dir / "downloaded" / "stock_prediction_model_anchored-path-v1.keras"
            downloaded_path.parent.mkdir(parents=True)
            downloaded_path.write_text("model", encoding="utf-8")
            settings = Settings(
                app_name="BuckLakeAI",
                app_version="2.0.0",
                model_version="anchored-path-v1",
                input_contract_version="v2.1.0",
                root_dir=root_dir,
                hf_repo_id="test/repo",
                hf_model_filename="stock_prediction_model_anchored-path-v1.keras",
                hf_scaler_filename="scalers_v21_rolling_7d.json",
                hf_cache_dir=cache_dir,
                hf_token="token-1",
                model_weights_path=root_dir / "missing.keras",
                text_encoder_model="test-encoder",
                text_encoder_device="cpu",
                text_encoder_max_seq_length=512,
                enable_finbert_sentiment=False,
                finbert_model_name="ProsusAI/finbert",
                preload_model=False,
                preload_text_encoder=False,
                scaler_artifact_path=root_dir / "scalers.json",
            )

            with patch(
                "bucklake_ai.model_loader.hf_hub_download",
                return_value=str(downloaded_path),
            ) as download:
                resolved = ModelLoader(settings)._resolve_model_path()

            self.assertEqual(resolved, downloaded_path)
            download.assert_called_once_with(
                repo_id="test/repo",
                filename="stock_prediction_model_anchored-path-v1.keras",
                cache_dir=str(cache_dir),
                token="token-1",
            )

    def test_anchored_path_rejects_input_contract_mismatch(self):
        model = StubModel()
        model.inputs = [StubTensor("wrong_input", (None, 768))] + model.inputs[1:]

        with self.assertRaisesRegex(ModelArtifactNotFound, "input contract mismatch"):
            ModelLoader._validate_anchored_path_model(model)

    def test_anchored_path_v2_accepts_index_path_output_contract(self):
        ModelLoader._validate_anchored_path_model(StubModelV2(), "anchored-path-v2")

    def test_anchored_path_v2_rejects_missing_index_outputs(self):
        with self.assertRaisesRegex(ModelArtifactNotFound, "anchored-path-v2 artifact output contract mismatch"):
            ModelLoader._validate_anchored_path_model(StubModel(), "anchored-path-v2")

    def test_anchored_path_v2_rejects_index_output_shape_mismatch(self):
        model = StubModelV2()
        model.outputs = list(model.outputs)
        model.outputs[8] = StubOutput((None, 3, 1))

        with self.assertRaisesRegex(ModelArtifactNotFound, "output shape mismatch"):
            ModelLoader._validate_anchored_path_model(model, "anchored-path-v2")

    def test_anchored_path_v1_rejects_stock_output_shape_mismatch(self):
        model = StubModel()
        model.outputs = [StubOutput((None, 3, 1)) for _ in range(8)]
        model.outputs[4] = StubOutput((None, 3, 1))

        with self.assertRaisesRegex(ModelArtifactNotFound, "output shape mismatch"):
            ModelLoader._validate_anchored_path_model(model)

    def test_anchored_path_rejects_artifact_saved_by_newer_keras(self):
        with tempfile.TemporaryDirectory() as tmp:
            cache_dir = Path(tmp) / ".cache" / "huggingface"
            cache_dir.mkdir(parents=True)
            artifact_path = cache_dir / "anchored_path_v1.keras"
            with zipfile.ZipFile(artifact_path, "w") as archive:
                archive.writestr("metadata.json", '{"keras_version": "99.0.0"}')

            settings = Settings(
                app_name="BuckLakeAI",
                app_version="2.0.0",
                model_version="anchored-path-v1",
                input_contract_version="v2.1.0",
                root_dir=Path(tmp),
                hf_repo_id="parkerjj/BuckLake-Stock-Model",
                hf_model_filename="anchored_path_v1.keras",
                hf_scaler_filename="scalers_v21_rolling_7d.json",
                hf_cache_dir=cache_dir,
                hf_token=None,
                model_weights_path=artifact_path,
                text_encoder_model="test-encoder",
                text_encoder_device="cpu",
                text_encoder_max_seq_length=512,
                enable_finbert_sentiment=False,
                finbert_model_name="ProsusAI/finbert",
                preload_model=False,
                preload_text_encoder=False,
                scaler_artifact_path=Path(tmp) / "scalers.json",
            )

            with self.assertRaisesRegex(ModelArtifactNotFound, "saved with Keras 99.0.0"):
                ModelLoader(settings).get_model()


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