MLModels / tests /test_pipeline.py
pcsekhar's picture
Deploy house price predictor
9e03a51 verified
Raw
History Blame Contribute Delete
1.98 kB
from pathlib import Path
import pandas as pd
from src.modeling import HousePriceModel, train_model
def sample_training_frame() -> pd.DataFrame:
return pd.DataFrame(
{
"OverallQual": [5, 6, 7, 8, 9, 4, 6, 7],
"GrLivArea": [1200, 1500, 1800, 2100, 2400, 900, 1600, 2000],
"GarageCars": [1, 2, 2, 3, 3, 1, 2, 2],
"TotalBsmtSF": [800, 900, 1100, 1300, 1500, 600, 950, 1250],
"FullBath": [1, 2, 2, 2, 3, 1, 2, 2],
"YearBuilt": [1960, 1975, 1990, 2001, 2010, 1950, 1982, 1998],
"Neighborhood": ["NAmes", "CollgCr", "Somerst", "NridgHt", "NoRidge", "OldTown", "NAmes", "Gilbert"],
"HouseStyle": ["1Story", "1Story", "2Story", "2Story", "2Story", "1Story", "1Story", "2Story"],
"SalePrice": [135000, 165000, 210000, 260000, 320000, 110000, 175000, 235000],
}
)
def test_train_model_persists_artifacts_and_reports_metrics(tmp_path: Path) -> None:
artifact_path = tmp_path / "house_price_model.joblib"
trained = train_model(sample_training_frame(), artifact_path=artifact_path)
assert artifact_path.exists()
assert trained.metrics["rmse"] >= 0
assert trained.metrics["r2"] <= 1
assert "OverallQual" in trained.feature_names
assert "SalePrice" not in trained.feature_names
def test_house_price_model_predicts_positive_value(tmp_path: Path) -> None:
artifact_path = tmp_path / "house_price_model.joblib"
trained = train_model(sample_training_frame(), artifact_path=artifact_path)
model = HousePriceModel.load(artifact_path)
prediction = model.predict(
{
"OverallQual": 7,
"GrLivArea": 1850,
"GarageCars": 2,
"TotalBsmtSF": 1050,
"FullBath": 2,
"YearBuilt": 1995,
"Neighborhood": "Somerst",
"HouseStyle": "2Story",
}
)
assert prediction > 0
assert model.metrics == trained.metrics