Pythonでテストを書く習慣【pytestでMLコードの品質を保つ方法】

はじめに

機械学習のコードは「動けばいい」で書かれがちですが、チーム開発や長期運用では自動テストが欠かせません。「コードを変更したら精度が下がった」「データ前処理のバグで本番が壊れた」という問題をテストで防ぐ方法を解説します。

pytestの基本

pip install pytest

# test_features.py
import pytest
import pandas as pd
import numpy as np
from src.features import create_lag_features, encode_categories

def test_lag_features_basic():
    """ラグ特徴量が正しく作成されることを確認"""
    df = pd.DataFrame({'value': [1, 2, 3, 4, 5]})
    result = create_lag_features(df, 'value', lags=[1, 2])
    
    assert 'value_lag1' in result.columns
    assert 'value_lag2' in result.columns
    assert result['value_lag1'].iloc[1] == 1  # 1行目のlag1は元の0行目の値

def test_lag_features_with_nulls():
    """ラグ特徴量の先頭行はNaNになることを確認"""
    df = pd.DataFrame({'value': [10, 20, 30]})
    result = create_lag_features(df, 'value', lags=[1])
    
    assert pd.isna(result['value_lag1'].iloc[0])  # 先頭はNaN
    assert result['value_lag1'].iloc[1] == 10

def test_encode_categories():
    """カテゴリエンコーディングが正しく動作することを確認"""
    df = pd.DataFrame({'category': ['A', 'B', 'A', 'C']})
    result = encode_categories(df, 'category')
    
    assert result['category'].dtype in [np.int64, np.int32]
    assert result['category'].nunique() == 3  # A, B, C の3種類

@pytest.fixture
def sample_df():
    """共通のサンプルデータ（複数のテストで再利用）"""
    return pd.DataFrame({
        'feature1': [1.0, 2.0, np.nan, 4.0],
        'feature2': [10, 20, 30, 40],
        'target': [0, 1, 1, 0]
    })

def test_no_missing_values_after_preprocessing(sample_df):
    """前処理後に欠損値がないことを確認"""
    from src.features import preprocess
    result = preprocess(sample_df)
    assert result.isnull().sum().sum() == 0  # 欠損値なし

# テストの実行
pytest tests/ -v

# カバレッジの確認
pytest tests/ --cov=src --cov-report=html

モデルのテスト

# test_model.py
def test_model_prediction_shape(trained_model, sample_features):
    """予測結果の形状が正しいことを確認"""
    predictions = trained_model.predict(sample_features)
    assert predictions.shape == (len(sample_features),)

def test_model_probability_range(trained_model, sample_features):
    """確率値が0〜1の範囲内であることを確認"""
    probabilities = trained_model.predict_proba(sample_features)
    assert (probabilities >= 0).all() and (probabilities <= 1).all()
    assert np.allclose(probabilities.sum(axis=1), 1.0)  # 合計が1

def test_model_performance(trained_model, test_data):
    """テストデータでの性能が閾値以上であることを確認"""
    from sklearn.metrics import roc_auc_score
    X_test, y_test = test_data
    y_proba = trained_model.predict_proba(X_test)[:, 1]
    auc = roc_auc_score(y_test, y_proba)
    assert auc >= 0.75, f'AUCが閾値を下回りました: {auc:.3f} < 0.75'

まとめ

機械学習コードのテストは「前処理関数の単体テスト」「モデルの出力形状・範囲のテスト」「性能が閾値を下回ったら失敗するテスト」の3種類から始めましょう。pytestとGitHub Actionsを組み合わせると、コードをpushするたびに自動テストが走るCI/CDパイプラインが構築できます。転職ポートフォリオにテストコードがあると、コード品質への意識が高い印象を与えられます。