File size: 5,764 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 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 | """
Integration test: DataAgent → FeatureEngineering → Prediction
Demonstrates complete flow from natural language query to win probability.
Run with:
python -m pytest tests/test_integration_agent_predict.py -v -s
"""
import json
import os
from pathlib import Path
import pandas as pd
import pytest
from agents.data_agent import data_agent, QueryIntent
from ml.feature_engineering import prepare_model_data
from ml.predict import predict_dataframe, load_model_and_encoders
@pytest.fixture
def minimal_model_and_encoders():
"""Load real trained model from MLflow."""
try:
run_id = os.getenv("KRONECTOR_TEST_RUN_ID")
if not run_id:
pytest.skip("KRONECTOR_TEST_RUN_ID not set")
model, encoders = load_model_and_encoders(run_id)
return model, encoders
except Exception as e:
pytest.skip(f"Model not available: {e}")
def test_agent_to_prediction_pipeline():
"""
Test: Natural language → Intent → DataFrame → Predictions
Uses mock parser (no API key needed) to verify the full pipeline.
"""
# Mock parser for reproducibility
def mock_parser(query: str) -> QueryIntent:
"""For this test, always return 2023 Bahrain Verstappen."""
return {
"season": 2023,
"round": 1,
"driver_id": "VER",
"driver_name": "Max Verstappen",
}
# Step 1: Natural language to structured data
query = "Will Max Verstappen win at Bahrain 2023?"
result = data_agent(query, parser=mock_parser)
# Verify intent extraction
assert result["intent"]["season"] == 2023
assert result["intent"]["round"] == 1
assert result["intent"]["driver_id"] == "VER"
print(f"✓ Intent parsed: {result['intent']}")
# Step 2: Verify DataFrame is prediction-compatible
df = result["dataframe"]
assert len(df) == 1
assert "driver_id" in df.columns
assert df.iloc[0]["driver_id"] == "VER"
print(f"✓ DataFrame shape: {df.shape}")
print(f"✓ Columns: {list(df.columns)[:5]}... ({len(df.columns)} total)")
# Step 3: Feature engineering (no encoders needed for shape test)
bundle, encoders = prepare_model_data(df)
assert len(bundle.X) == 1
assert len(bundle.feature_columns) > 0
assert "grid_position" in bundle.feature_columns
assert "win_probability" not in bundle.feature_columns
print(f"✓ Features: {len(bundle.feature_columns)} columns, shape {bundle.X.shape}")
# Step 4: Can generate predictions (if model available)
try:
run_id = os.getenv("KRONECTOR_TEST_RUN_ID")
if not run_id:
print("ℹ Model inference skipped: KRONECTOR_TEST_RUN_ID not set")
return
model, fitted_encoders = load_model_and_encoders(run_id)
predictions = predict_dataframe(df, model, fitted_encoders)
win_prob = predictions.iloc[0]["win_probability"]
print(f"✓ Prediction: {win_prob:.2%} win probability for Verstappen")
assert 0 <= win_prob <= 1
except Exception as e:
print(f"ℹ Model inference skipped: {e}")
print(" (This is expected if MLflow model is not available)")
def test_multiple_drivers_same_race():
"""Test querying all drivers in a race without driver filter."""
def mock_parser(query: str) -> QueryIntent:
"""Return race without driver filter."""
return {
"season": 2023,
"round": 1,
"driver_id": None, # Get all drivers
"driver_name": None,
}
result = data_agent("Predict Bahrain 2023 standings", parser=mock_parser)
# Should have multiple drivers
num_drivers = len(result["dataframe"])
assert num_drivers > 1
print(f"✓ Queried {num_drivers} drivers for 2023 Bahrain")
# All rows should have different driver_ids
driver_ids = result["dataframe"]["driver_id"].unique()
assert len(driver_ids) == num_drivers
print(f"✓ Unique drivers: {', '.join(sorted(driver_ids))}")
def test_agent_output_matches_prediction_schema():
"""Verify agent output matches PredictionInputRow schema."""
def mock_parser(query: str) -> QueryIntent:
return {"season": 2023, "round": 1, "driver_id": "HAM"}
result = data_agent("Hamilton Bahrain", parser=mock_parser)
rows = result["rows"]
# Verify required fields exist
required_fields = {
"season", "round", "driver_id", "team",
"grid_position", "finish_position", "circuit_id"
}
for row in rows:
for field in required_fields:
assert field in row, f"Missing field: {field}"
print(f"✓ All {len(required_fields)} required fields present")
print(f"✓ Sample row keys: {list(rows[0].keys())}")
def test_agent_error_handling():
"""Test agent gracefully handles invalid queries."""
def mock_parser(query: str) -> QueryIntent:
# Return an impossible race (1000 rounds doesn't exist)
return {"season": 2099, "round": 1000, "driver_id": None}
with pytest.raises(ValueError, match="No rows found"):
data_agent("Impossible race", parser=mock_parser)
print("✓ Agent properly rejects non-existent races")
if __name__ == "__main__":
# Run standalone (without pytest)
print("=" * 60)
print("KRONECTOR: DataAgent Integration Test")
print("=" * 60)
test_agent_to_prediction_pipeline()
print()
test_multiple_drivers_same_race()
print()
test_agent_output_matches_prediction_schema()
print()
test_agent_error_handling()
print()
print("=" * 60)
print("All integration tests passed!")
print("=" * 60)
|