tests/test_model.py - master - nevers/hackINTERPOL

tests/test_model.py
101 строка4 KB
nevers
upload files
13 янв 2026, 17:14
13 янв 2026, 17:14d476ef2
"""Tests for risk prediction model."""

import pandas as pd
import pytest

from studypulse.metrics import compute_risk_flags
from studypulse.model import predict_risk, train_and_predict, train_risk_model


@pytest.fixture
def sample_df() -> pd.DataFrame:
    """Create sample DataFrame for testing."""
    return pd.DataFrame(
        {
            "student_id": ["STU001", "STU002", "STU003", "STU004", "STU005"],
            "student_name": ["Alice", "Bob", "Charlie", "Diana", "Edward"],
            "assignment_score": [85.0, 45.0, 95.0, 72.0, 58.0],
            "assignments_submitted": [8, 4, 9, 7, 5],
            "attendance_rate": [0.9, 0.6, 0.98, 0.75, 0.68],
            "days_since_last_submission": [3, 18, 1, 12, 15],
            "days_to_deadline": [5, -3, 7, 2, -1],
        }
    )


def test_predict_risk_with_heuristic_fallback(sample_df: pd.DataFrame) -> None:
    """Test risk prediction uses heuristic fallback when no model provided."""
    df_with_flags = compute_risk_flags(sample_df)
    result = predict_risk(df_with_flags, model=None)

    assert "predicted_risk_flag" in result.columns
    assert "predicted_risk_probability" in result.columns

    # High-risk student (STU002) should be flagged
    stu002_risk = result.loc[result["student_id"] == "STU002", "predicted_risk_flag"].iloc[0]
    assert stu002_risk == 1

    # Low-risk student (STU003) should not be flagged
    stu003_risk = result.loc[result["student_id"] == "STU003", "predicted_risk_flag"].iloc[0]
    assert stu003_risk == 0


def test_predict_risk_with_model(sample_df: pd.DataFrame) -> None:
    """Test risk prediction with trained model."""
    df_with_flags = compute_risk_flags(sample_df)
    model = train_risk_model(df_with_flags)
    result = predict_risk(df_with_flags, model=model)

    assert "predicted_risk_flag" in result.columns
    assert "predicted_risk_probability" in result.columns
    # Probabilities should be between 0 and 1
    assert result["predicted_risk_probability"].min() >= 0.0
    assert result["predicted_risk_probability"].max() <= 1.0


def test_train_and_predict_small_dataset() -> None:
    """Test train_and_predict falls back to heuristic for small datasets."""
    small_df = pd.DataFrame(
        {
            "student_id": ["STU001", "STU002"],
            "student_name": ["Alice", "Bob"],
            "assignment_score": [85.0, 45.0],
            "assignments_submitted": [8, 4],
            "attendance_rate": [0.9, 0.6],
            "days_since_last_submission": [3, 18],
            "days_to_deadline": [5, -3],
        }
    )

    result = train_and_predict(small_df)
    assert "predicted_risk_flag" in result.columns
    assert "predicted_risk_probability" in result.columns
    # Should use heuristic fallback for small dataset
    assert result.loc[result["student_id"] == "STU002", "predicted_risk_flag"].iloc[0] == 1


def test_train_and_predict_large_dataset(sample_df: pd.DataFrame) -> None:
    """Test train_and_predict uses model for larger datasets."""
    # Create a larger dataset by duplicating
    large_df = pd.concat([sample_df] * 3, ignore_index=True)
    large_df["student_id"] = [f"STU{i:03d}" for i in range(len(large_df))]

    result = train_and_predict(large_df)
    assert "predicted_risk_flag" in result.columns
    assert "predicted_risk_probability" in result.columns
    # Should train and use model for larger dataset
    assert len(result) == len(large_df)


def test_model_deterministic_behavior(sample_df: pd.DataFrame) -> None:
    """Test that model training produces deterministic results."""
    df_with_flags = compute_risk_flags(sample_df)

    model1 = train_risk_model(df_with_flags)
    model2 = train_risk_model(df_with_flags)

    result1 = predict_risk(df_with_flags, model=model1)
    result2 = predict_risk(df_with_flags, model=model2)

    # Predictions should be identical due to fixed random seed
    pd.testing.assert_frame_equal(result1, result2)