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()