File size: 7,260 Bytes
fcf8749 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | """
Tests for DriverEffortLearner XGBoost implementation.
"""
import pytest
import numpy as np
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import uuid4
from datetime import datetime, timedelta
import sys
sys.path.insert(0, str(__file__).replace("\\", "/").rsplit("/tests", 1)[0])
from app.services.learning_agent import DriverEffortLearner
class TestDriverEffortLearner:
"""Tests for DriverEffortLearner class."""
@pytest.fixture
def mock_db(self):
"""Create mock async database session."""
mock = AsyncMock()
mock.execute = AsyncMock()
mock.add = MagicMock()
return mock
@pytest.fixture
def learner(self, mock_db):
"""Create DriverEffortLearner instance."""
return DriverEffortLearner(mock_db)
def test_feature_names_defined(self, learner):
"""Test that feature names are properly defined."""
assert len(learner.FEATURE_NAMES) > 0
assert "num_packages" in learner.FEATURE_NAMES
assert "total_weight_kg" in learner.FEATURE_NAMES
assert "num_stops" in learner.FEATURE_NAMES
def test_min_training_samples(self, learner):
"""Test minimum training samples requirement."""
assert learner.MIN_TRAINING_SAMPLES == 10
assert learner.MAX_TRAINING_SAMPLES == 100
@pytest.mark.asyncio
async def test_load_model_no_record(self, learner, mock_db):
"""Test load_model returns None when no record exists."""
mock_db.execute.return_value = MagicMock(
scalar_one_or_none=MagicMock(return_value=None)
)
driver_id = uuid4()
model = await learner.load_model(driver_id)
assert model is None
@pytest.mark.asyncio
async def test_get_model_version_no_record(self, learner, mock_db):
"""Test get_model_version returns None when no record exists."""
mock_db.execute.return_value = MagicMock(
scalar_one_or_none=MagicMock(return_value=None)
)
driver_id = uuid4()
version = await learner.get_model_version(driver_id)
assert version is None
@pytest.mark.asyncio
async def test_predict_effort_no_model(self, learner, mock_db):
"""Test predict_effort returns None when no model exists."""
mock_db.execute.return_value = MagicMock(
scalar_one_or_none=MagicMock(return_value=None)
)
driver_id = uuid4()
route_features = {
"num_packages": 20,
"total_weight_kg": 50.0,
"num_stops": 10,
}
prediction, version = await learner.predict_effort(driver_id, route_features)
assert prediction is None
assert version is None
@pytest.mark.asyncio
async def test_update_model_insufficient_data(self, learner, mock_db):
"""Test update_model skips when insufficient data."""
# Return only 5 records (less than MIN_TRAINING_SAMPLES)
mock_result = MagicMock()
mock_result.scalars.return_value.all.return_value = [
MagicMock(actual_effort=50.0) for _ in range(5)
]
mock_db.execute.return_value = mock_result
driver_id = uuid4()
result = await learner.update_model(driver_id)
assert result["status"] == "skipped"
assert result["reason"] == "insufficient_data"
assert result["samples"] == 5
@pytest.mark.asyncio
async def test_get_model_status_no_record(self, learner, mock_db):
"""Test get_model_status returns None when no record exists."""
mock_db.execute.return_value = MagicMock(
scalar_one_or_none=MagicMock(return_value=None)
)
driver_id = uuid4()
status = await learner.get_model_status(driver_id)
assert status is None
@pytest.mark.asyncio
async def test_get_model_status_with_record(self, learner, mock_db):
"""Test get_model_status returns correct data."""
driver_id = uuid4()
mock_model = MagicMock(
model_version=3,
training_samples=50,
current_mse=8.5,
r2_score=0.85,
mse_history={"values": [10.0, 9.0, 8.5]},
active=True,
last_trained_at=datetime.utcnow(),
)
mock_db.execute.return_value = MagicMock(
scalar_one_or_none=MagicMock(return_value=mock_model)
)
status = await learner.get_model_status(driver_id)
assert status is not None
assert status["model_version"] == 3
assert status["training_samples"] == 50
assert status["current_mse"] == 8.5
assert status["r2_score"] == 0.85
assert status["active"] is True
class TestXGBoostIntegration:
"""Tests for XGBoost integration (requires xgboost installed)."""
@pytest.fixture
def mock_db(self):
mock = AsyncMock()
mock.execute = AsyncMock()
mock.add = MagicMock()
return mock
def test_xgboost_availability_check(self, mock_db):
"""Test that XGBoost availability is properly detected."""
learner = DriverEffortLearner(mock_db)
# Should be True if xgboost is installed, False otherwise
assert isinstance(learner._xgb_available, bool)
@pytest.mark.skipif(
not DriverEffortLearner(AsyncMock())._xgb_available,
reason="XGBoost not installed"
)
def test_xgboost_can_train(self):
"""Test that XGBoost model can be trained manually."""
import xgboost as xgb
import pandas as pd
# Create synthetic data
X = pd.DataFrame({
"num_packages": [10, 20, 30, 15, 25],
"total_weight_kg": [30, 50, 80, 40, 60],
"num_stops": [5, 10, 15, 7, 12],
})
y = np.array([40, 60, 85, 50, 70])
model = xgb.XGBRegressor(n_estimators=10, max_depth=3)
model.fit(X, y)
predictions = model.predict(X)
assert len(predictions) == 5
assert all(isinstance(p, (float, np.floating)) for p in predictions)
class TestModelSerialization:
"""Tests for model serialization."""
@pytest.mark.skipif(
not DriverEffortLearner(AsyncMock())._xgb_available,
reason="XGBoost not installed"
)
def test_model_pickle_roundtrip(self):
"""Test that model can be pickled and unpickled."""
import pickle
import xgboost as xgb
import pandas as pd
# Train a small model
X = pd.DataFrame({
"num_packages": [10, 20, 30],
"total_weight_kg": [30, 50, 80],
})
y = np.array([40, 60, 85])
model = xgb.XGBRegressor(n_estimators=5)
model.fit(X, y)
# Pickle and unpickle
pickled = pickle.dumps(model)
loaded = pickle.loads(pickled)
# Should produce same predictions
original_preds = model.predict(X)
loaded_preds = loaded.predict(X)
np.testing.assert_array_almost_equal(original_preds, loaded_preds)
|