| """ |
| 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.""" |
| |
| 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) |
| |
| |
| 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 |
| |
| |
| 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 |
| |
| |
| 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) |
| |
| |
| pickled = pickle.dumps(model) |
| loaded = pickle.loads(pickled) |
| |
| |
| original_preds = model.predict(X) |
| loaded_preds = loaded.predict(X) |
| |
| np.testing.assert_array_almost_equal(original_preds, loaded_preds) |
|
|
