Predictive Testing Agent

from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.pipeline import Pipeline
import joblib

class TestFailurePredictor:
    def __init__(self):
        self.pipeline = Pipeline([
            ('tfidf', TfidfVectorizer(max_features=1000)),
            ('clf', RandomForestClassifier(
                n_estimators=100,
                class_weight='balanced',
            )),
        ])

    def train(self, historical_data: pd.DataFrame):
        X = historical_data['features']
        y = historical_data['failed']
        self.pipeline.fit(X, y)

    def predict(self, change_features: list[str]) -> list[float]:
        probabilities = self.pipeline.predict_proba(
            change_features
        )
        return [p[1] for p in probabilities]

    def select_tests(
        self,
        tests: list[Test],
        threshold: float = 0.3,
        time_budget: int = 600,
    ) -> list[Test]:
        features = [t.feature_vector for t in tests]
        scores = self.predict(features)
        ranked = sorted(
            zip(tests, scores),
            key=lambda x: -x[1],
        )
        selected = []
        total_time = 0
        for test, score in ranked:
            if score < threshold:
                break
            if total_time + test.est_duration > time_budget:
                break
            selected.append(test)
            total_time += test.est_duration
        return selected

from predictive_agent.features import FeatureExtractor

extractor = FeatureExtractor()

pr_data = {
    "files": [
        "src/api/checkout.py",
        "tests/test_checkout.py",
    ],
    "author": "ervin",
    "commit_message": "fix: handle empty cart",
    "changed_lines": {"src/api/checkout.py": [45, 67]},
}

feature_vector = extractor.extract(pr_data)

# Feature components:
# - File path TF-IDF tokens
# - Author one-hot encoding
# - Commit message tokens
# - Coverage mapping (which tests touch changed files)
# - Historical failure rate per test

print(f"Feature dimensions: {feature_vector.shape}")
# Output: Feature dimensions: (1, 1000)

# Prediction per test
predictor = TestFailurePredictor.load("models/v1.pkl")
tests = get_all_tests()

scores = predictor.predict([
    feature_vector for _ in tests
])

for test, score in sorted(
    zip(tests, scores),
    key=lambda x: -x[1],
)[:5]:
    print(f"{test.name}: {score:.1%} failure risk")