File size: 3,604 Bytes
2532605 | 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 | """
Tests for agents.data_agent.
Run: python -m pytest tests/test_data_agent.py -v
"""
from types import SimpleNamespace
import pandas as pd
import pytest
from agents.data_agent import (
data_agent,
parse_query_with_groq,
)
from ml.feature_engineering import prepare_model_data
@pytest.fixture
def agent_dataset(tmp_path):
df = pd.DataFrame(
{
"season": [2023, 2023],
"round": [1, 1],
"driver_id": ["VER", "HAM"],
"driver_name": ["Max Verstappen", "Lewis Hamilton"],
"team": ["Red Bull Racing", "Mercedes"],
"grid_position": [1.0, 2.0],
"finish_position": [1.0, 2.0],
"circuit_id": ["Bahrain Grand Prix", "Bahrain Grand Prix"],
"sector_1_time": [28.7, 28.9],
"sector_2_time": [38.5, 38.7],
"sector_3_time": [22.4, 22.6],
"avg_lap_time_practice": [95.1, 95.4],
"tire_compound": [0.0, 1.0],
"tire_age_laps": [14.0, 12.0],
"fresh_tire": [0.0, 1.0],
"pit_stop_count": [2.0, 2.0],
"team_pit_speed": [2.5, 2.8],
"weather_temp_track": [31.0, 31.0],
"weather_rainfall": [0.0, 0.0],
"telemetry_available": [True, True],
}
)
path = tmp_path / "races.parquet"
df.to_parquet(path, index=False)
return path
def test_data_agent_returns_prediction_compatible_dataframe(agent_dataset):
def parser(query):
return {
"season": 2023,
"round": 1,
"driver_id": "VER",
"driver_name": "Max Verstappen",
}
result = data_agent("Will Verstappen win Bahrain 2023?", agent_dataset, parser)
assert result["intent"]["driver_id"] == "VER"
assert len(result["rows"]) == 1
assert result["dataframe"].iloc[0]["driver_id"] == "VER"
bundle, _ = prepare_model_data(result["dataframe"])
assert len(bundle.X) == 1
def test_data_agent_returns_all_race_rows_when_no_driver(agent_dataset):
def parser(query):
return {"season": 2023, "round": 1, "driver_id": None, "driver_name": None}
result = data_agent("Predict Bahrain 2023", agent_dataset, parser)
assert len(result["dataframe"]) == 2
def test_data_agent_raises_for_missing_race(agent_dataset):
def parser(query):
return {"season": 2024, "round": 99, "driver_id": None, "driver_name": None}
with pytest.raises(ValueError, match="No rows found"):
data_agent("Predict a missing race", agent_dataset, parser)
def test_parse_query_with_groq_parses_json(monkeypatch):
class FakeCompletions:
def create(self, **kwargs):
return SimpleNamespace(
choices=[
SimpleNamespace(
message=SimpleNamespace(
content=(
'{"season": 2023, "round": 1, '
'"driver_id": "ver", "driver_name": "Max Verstappen"}'
)
)
)
]
)
class FakeGroq:
def __init__(self, api_key):
self.chat = SimpleNamespace(
completions=FakeCompletions()
)
monkeypatch.setenv("GROQ_API_KEY", "test-key")
monkeypatch.setattr("agents.data_agent.Groq", FakeGroq)
intent = parse_query_with_groq("Will Max win Bahrain 2023?")
assert intent == {
"season": 2023,
"round": 1,
"driver_id": "VER",
"driver_name": "Max Verstappen",
}
|