Refactor Python core with explicit contracts and stronger test coverage

Author: oceanicpatterns

ml_vix_model.py

@@ -1,6 +1,10 @@
+from __future__ import annotations
+
 import os
 import re
+from dataclasses import dataclass
 from io import StringIO
+from typing import Mapping, Sequence
 
 import pandas as pd
 import requests
@@ -16,6 +20,22 @@
 TABLE_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*){2}$")
 
 
+@dataclass(frozen=True)
+class TrainingResult:
+    mse: float
+    train_rows: int
+    test_rows: int
+    predicted_close_for_input: float
+
+    def to_dict(self) -> dict[str, float | int]:
+        return {
+            "mse": self.mse,
+            "train_rows": self.train_rows,
+            "test_rows": self.test_rows,
+            "predicted_close_for_input": self.predicted_close_for_input,
+        }
+
+
 class DataProcessor:
     def __init__(self, url: str):
         self.url = url
@@ -47,7 +67,7 @@ def _validate_table_name(table_name: str) -> str:
     return table_name
 
 
-def create_temp_table(conn, table_name: str) -> None:
+def create_temp_table(conn: object, table_name: str) -> None:
     table = _validate_table_name(table_name)
     with conn.cursor() as cur:
         cur.execute(f"DROP TABLE IF EXISTS {table}")
@@ -57,7 +77,7 @@ def create_temp_table(conn, table_name: str) -> None:
         )
 
 
-def insert_data_to_temp_table(conn, df: pd.DataFrame, table_name: str) -> None:
+def insert_data_to_temp_table(conn: object, df: pd.DataFrame, table_name: str) -> None:
     table = _validate_table_name(table_name)
     with conn.cursor() as cur:
         payload = df.copy()
@@ -66,16 +86,53 @@ def insert_data_to_temp_table(conn, df: pd.DataFrame, table_name: str) -> None:
         cur.executemany(f"INSERT INTO {table} VALUES (%s, %s, %s)", rows)
 
 
-def fetch_data_from_temp_table(conn, table_name: str) -> pd.DataFrame:
+def fetch_data_from_temp_table(conn: object, table_name: str) -> pd.DataFrame:
     table = _validate_table_name(table_name)
     with conn.cursor() as cur:
         cur.execute(f"SELECT * FROM {table}")
         rows = cur.fetchall()
     return pd.DataFrame(rows, columns=["DATE", "CLOSE_PRICE", "VOLATILITY_INDEX"])
 
 
-def run_training_and_generate_results() -> dict:
+def _fit_and_evaluate(
+    features: pd.DataFrame, targets: pd.Series, predict_for: float
+) -> TrainingResult:
+    if features.empty or targets.empty:
+        raise ValueError("Training data is empty. Cannot fit model.")
+
+    X_train, X_test, y_train, y_test = train_test_split(
+        features, targets, test_size=0.2, random_state=42
+    )
+    if X_train.empty or X_test.empty:
+        raise ValueError("Insufficient rows after train-test split.")
+
+    model = LinearRegression()
+    model.fit(X_train, y_train)
+
+    predictions = model.predict(X_test)
+    mse = mean_squared_error(y_test, predictions)
+    single_prediction = float(model.predict(pd.DataFrame({"VOLATILITY_INDEX": [predict_for]}))[0])
+
+    return TrainingResult(
+        mse=float(mse),
+        train_rows=int(len(X_train)),
+        test_rows=int(len(X_test)),
+        predicted_close_for_input=single_prediction,
+    )
+
+
+def _build_report(result: TrainingResult) -> str:
+    return (
+        "Our model predicts VIX CLOSE_PRICE from VOLATILITY_INDEX. "
+        f"MSE on unseen data: {result.mse:.4f}. "
+        f"Train rows: {result.train_rows}, Test rows: {result.test_rows}. "
+        f"Sample prediction @ VOLATILITY_INDEX=0.40: {result.predicted_close_for_input:.2f}."
+    )
+
+
+def run_training_and_generate_results() -> Mapping[str, float | int]:
     table_name = os.getenv("SNOWFLAKE_TEMP_TABLE", DEFAULT_TABLE)
+    prediction_input = float(os.getenv("VIX_PREDICTION_INPUT", "0.40"))
     conn = get_snowflake_connection()
     try:
         data_processor = DataProcessor(VIX_DATA_URL)
@@ -86,23 +143,13 @@ def run_training_and_generate_results() -> dict:
         insert_data_to_temp_table(conn, df, table_name=table_name)
         fetched_df = fetch_data_from_temp_table(conn, table_name=table_name)
 
-        X, y = fetched_df[["VOLATILITY_INDEX"]], fetched_df["CLOSE_PRICE"]
-        X_train, X_test, y_train, y_test = train_test_split(
-            X, y, test_size=0.2, random_state=42
-        )
-
-        model = LinearRegression()
-        model.fit(X_train, y_train)
-
-        predictions = model.predict(X_test)
-        mse = mean_squared_error(y_test, predictions)
-
-        report = (
-            "Our model predicts VIX CLOSE_PRICE from VOLATILITY_INDEX. "
-            f"The Mean Squared Error (MSE) on unseen test data is {mse:.4f}."
+        result = _fit_and_evaluate(
+            features=fetched_df[["VOLATILITY_INDEX"]],
+            targets=fetched_df["CLOSE_PRICE"],
+            predict_for=prediction_input,
         )
-        print(report)
-        return {"mse": float(mse), "train_rows": int(len(X_train)), "test_rows": int(len(X_test))}
+        print(_build_report(result))
+        return result.to_dict()
     finally:
         conn.close()
 

snowflake_connection.py

@@ -1,15 +1,39 @@
+from __future__ import annotations
+
 import configparser
 import os
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict
+from typing import Any, Dict, Mapping
 
 import snowflake.connector
+from snowflake.connector.connection import SnowflakeConnection
 
 
 REQUIRED_FIELDS = ("user", "password", "account", "warehouse", "database", "schema")
 
 
-def _load_from_config(config_path: Path) -> Dict[str, str]:
+@dataclass(frozen=True)
+class SnowflakeSettings:
+    user: str
+    password: str
+    account: str
+    warehouse: str
+    database: str
+    schema: str
+
+    def to_connect_kwargs(self) -> Dict[str, str]:
+        return {
+            "user": self.user,
+            "password": self.password,
+            "account": self.account,
+            "warehouse": self.warehouse,
+            "database": self.database,
+            "schema": self.schema,
+        }
+
+
+def _load_from_config(config_path: Path) -> Mapping[str, str]:
     parser = configparser.ConfigParser()
     if not config_path.exists():
         return {}
@@ -19,7 +43,7 @@ def _load_from_config(config_path: Path) -> Dict[str, str]:
     return {key: parser["snowflake"].get(key, "") for key in REQUIRED_FIELDS}
 
 
-def _load_from_env() -> Dict[str, str]:
+def _load_from_env() -> Mapping[str, str]:
     return {
         "user": os.getenv("SNOWFLAKE_USER", ""),
         "password": os.getenv("SNOWFLAKE_PASSWORD", ""),
@@ -30,7 +54,7 @@ def _load_from_env() -> Dict[str, str]:
     }
 
 
-def get_connection_params(config_path: str | None = None) -> Dict[str, str]:
+def get_connection_params(config_path: str | None = None) -> SnowflakeSettings:
     config_file = Path(
         config_path or os.getenv("SNOWFLAKE_CONFIG_PATH", "config/snowflake_config.ini")
     )
@@ -49,16 +73,11 @@ def get_connection_params(config_path: str | None = None) -> Dict[str, str]:
             f"values in {config_file}."
         )
 
-    return merged
+    return SnowflakeSettings(**merged)
 
 
-def get_snowflake_connection(config_path: str | None = None):
+def get_snowflake_connection(
+    config_path: str | None = None, **connector_kwargs: Any
+) -> SnowflakeConnection:
     params = get_connection_params(config_path=config_path)
-    return snowflake.connector.connect(
-        user=params["user"],
-        password=params["password"],
-        account=params["account"],
-        warehouse=params["warehouse"],
-        database=params["database"],
-        schema=params["schema"],
-    )
+    return snowflake.connector.connect(**params.to_connect_kwargs(), **connector_kwargs)

test_ml_vix_model.py

@@ -3,7 +3,7 @@
 import pandas as pd
 import pytest
 
-from ml_vix_model import DataProcessor, _validate_table_name
+from ml_vix_model import DataProcessor, _fit_and_evaluate, _validate_table_name
 
 
 def test_fetch_data_from_url_parses_required_columns():
@@ -42,3 +42,22 @@ def test_validate_table_name_accepts_fully_qualified_name():
 def test_validate_table_name_rejects_invalid_name():
     with pytest.raises(ValueError):
         _validate_table_name("raw.temp_table")
+
+
+def test_fit_and_evaluate_returns_expected_shape():
+    features = pd.DataFrame({"VOLATILITY_INDEX": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]})
+    targets = pd.Series([20, 25, 30, 35, 40, 45], name="CLOSE_PRICE")
+    result = _fit_and_evaluate(features=features, targets=targets, predict_for=0.4)
+    assert result.train_rows > 0
+    assert result.test_rows > 0
+    assert result.mse >= 0
+    assert isinstance(result.predicted_close_for_input, float)
+
+
+def test_fit_and_evaluate_raises_for_empty_data():
+    with pytest.raises(ValueError):
+        _fit_and_evaluate(
+            features=pd.DataFrame({"VOLATILITY_INDEX": []}),
+            targets=pd.Series([], dtype=float, name="CLOSE_PRICE"),
+            predict_for=0.4,
+        )

test_snowflake_connection.py

@@ -1,8 +1,50 @@
 import os
+from unittest.mock import patch
 
 import pytest
 
-from snowflake_connection import get_snowflake_connection
+from snowflake_connection import get_connection_params, get_snowflake_connection
+
+
+def test_get_connection_params_prefers_env(monkeypatch):
+    monkeypatch.setenv("SNOWFLAKE_USER", "u")
+    monkeypatch.setenv("SNOWFLAKE_PASSWORD", "p")
+    monkeypatch.setenv("SNOWFLAKE_ACCOUNT", "a")
+    monkeypatch.setenv("SNOWFLAKE_WAREHOUSE", "w")
+    monkeypatch.setenv("SNOWFLAKE_DATABASE", "d")
+    monkeypatch.setenv("SNOWFLAKE_SCHEMA", "s")
+    params = get_connection_params(config_path="config/does-not-exist.ini")
+    assert params.user == "u"
+    assert params.database == "d"
+
+
+def test_get_connection_params_missing_values_raises(monkeypatch):
+    for key in (
+        "SNOWFLAKE_USER",
+        "SNOWFLAKE_PASSWORD",
+        "SNOWFLAKE_ACCOUNT",
+        "SNOWFLAKE_WAREHOUSE",
+        "SNOWFLAKE_DATABASE",
+        "SNOWFLAKE_SCHEMA",
+    ):
+        monkeypatch.delenv(key, raising=False)
+
+    with pytest.raises(ValueError):
+        get_connection_params(config_path="config/does-not-exist.ini")
+
+
+def test_get_snowflake_connection_passes_connector_kwargs(monkeypatch):
+    monkeypatch.setenv("SNOWFLAKE_USER", "u")
+    monkeypatch.setenv("SNOWFLAKE_PASSWORD", "p")
+    monkeypatch.setenv("SNOWFLAKE_ACCOUNT", "a")
+    monkeypatch.setenv("SNOWFLAKE_WAREHOUSE", "w")
+    monkeypatch.setenv("SNOWFLAKE_DATABASE", "d")
+    monkeypatch.setenv("SNOWFLAKE_SCHEMA", "s")
+
+    with patch("snowflake_connection.snowflake.connector.connect") as connect_mock:
+        get_snowflake_connection(login_timeout=10)
+    connect_mock.assert_called_once()
+    assert connect_mock.call_args.kwargs["login_timeout"] == 10
 
 
 @pytest.mark.skipif(