tests: updated test files for model wrapper functionality

Author: W-Thurston

tests/test_linear_model_wrapper.py

@@ -0,0 +1,50 @@
+import numpy as np
+import pandas as pd
+
+from app.models.linear_model_wrapper import LinearModelWrapper
+
+
+def test_linear_wrapper_fit_and_predict():
+    """
+    Test that LinearModelWrapper can fit a model and return predictions.
+    Ensures the fitted model exposes .predict() and matches expected length.
+    """
+    X = pd.DataFrame({"x1": np.random.randn(100)})
+    y = 3 * X["x1"] + np.random.randn(100)
+
+    model = LinearModelWrapper(X, y).fit()
+    preds = model.predict()
+
+    assert len(preds) == len(y)
+    assert hasattr(model, "model")
+
+
+def test_linear_wrapper_residuals_and_fitted():
+    """
+    Verify that residuals + fitted values approximately equal the true target.
+    Confirms internal math and data shape integrity.
+    """
+    X = pd.DataFrame({"x1": np.random.randn(100)})
+    y = 2 * X["x1"] + np.random.randn(100)
+
+    model = LinearModelWrapper(X, y).fit()
+    residuals = model.residuals()
+    fitted = model.fitted()
+
+    # residuals = y - y_pred
+    np.testing.assert_allclose(y.values, residuals + fitted, rtol=1e-4)
+
+
+def test_linear_wrapper_summary():
+    """
+    Confirm the summary() method returns expected keys and types.
+    """
+    X = pd.DataFrame({"x1": np.random.randn(50)})
+    y = X["x1"] + np.random.randn(50)
+
+    model = LinearModelWrapper(X, y).fit()
+    summary = model.summary()
+
+    assert "model_type" in summary
+    assert summary["model_type"].lower() == "linear regression"
+    assert 0 <= summary["r_squared"] <= 1

tests/test_linearity.py

@@ -1,6 +1,11 @@
 # tests/test_linearity.py
+import numpy as np
+import pandas as pd
+
 from app.core import linearity
+from app.core.linearity import check_linearity
 from app.data import simulated_data
+from app.models.linear_model_wrapper import LinearModelWrapper
 
 
 def test_linearity_r_squared_threshold():
@@ -21,3 +26,15 @@ def test_linearity_plot_generation():
     result = linearity.check_linearity(df["x"], df["y"], return_plot=True)
     assert result.plot_base64 is not None
     assert result.plot_base64.startswith("iVBOR")  # PNG header in base64
+
+
+def test_linearity_with_model_wrapper():
+    """
+    Ensure check_linearity works when a pre-fit model_wrapper is provided.
+    """
+    X = pd.DataFrame({"x1": np.random.randn(100)})
+    y = 2 * X["x1"] + np.random.randn(100)
+    wrapper = LinearModelWrapper(X, y).fit()
+
+    result = check_linearity(X, y, model_wrapper=wrapper)
+    assert "r_squared" in result.details

tests/test_model_utils.py

@@ -0,0 +1,30 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from app.models.linear_model_wrapper import LinearModelWrapper
+from app.models.utils import get_model_wrapper
+
+
+def test_get_model_wrapper_linear():
+    """
+    Verify get_model_wrapper returns a LinearModelWrapper for 'linear' input.
+    """
+    X = pd.DataFrame({"x1": np.random.randn(30)})
+    y = 2 * X["x1"] + np.random.randn(30)
+
+    wrapper = get_model_wrapper("linear", X, y)
+    assert isinstance(wrapper, LinearModelWrapper)
+    assert hasattr(wrapper, "predict")
+    assert hasattr(wrapper, "residuals")
+
+
+def test_get_model_wrapper_invalid_type():
+    """
+    Confirm get_model_wrapper raises ValueError for unknown model_type.
+    """
+    X = pd.DataFrame({"x1": np.random.randn(30)})
+    y = 2 * X["x1"] + np.random.randn(30)
+
+    with pytest.raises(ValueError, match="Unsupported model type"):
+        get_model_wrapper("invalid_type", X, y)