defect-analysis/tests/test_ml_platform.py

import unittest

import pandas as pd

from defect_analysis.ml.datasets import build_supervised_dataset
from defect_analysis.ml.features import build_feature_frame
from defect_analysis.ml.image_models import ImageModelUnavailable, ImageModelWrapper
from defect_analysis.ml.model_registry import detect_optional_model_backends
from defect_analysis.ml.predict import predict_key_factors
from defect_analysis.ml.tabular_models import train_tabular_model
from defect_analysis.schemas import normalize_defect_schema


class MLPlatformTest(unittest.TestCase):
    def setUp(self):
        rows = []
        for i in range(40):
            hot = i < 24
            rows.append(
                {
                    "defect_id": f"D{i}",
                    "panel_id": f"P{i}",
                    "batch_id": "B1",
                    "equipment_id": "LAM-A01" if hot else "LAM-B01",
                    "seat_id": "R1C1" if hot else "R2C2",
                    "inspection_station": "AOI-1",
                    "timestamp": pd.Timestamp("2026-04-01 08:00:00"),
                    "defect_type": "气泡" if hot else "划痕",
                    "severity": "严重" if i % 5 == 0 else "轻微",
                    "x_mm": 10.0 + i,
                    "y_mm": 20.0,
                    "panel_width_mm": 155.0,
                    "panel_height_mm": 340.0,
                    "hour": 8,
                    "shift": "白班",
                    "day": "2026-04-01",
                    "lam_fixture_id": "FIX-HOT" if hot else "FIX-OK",
                    "material_lot_oca": "OCA-HOT" if hot else "OCA-OK",
                }
            )
        self.df = normalize_defect_schema(pd.DataFrame(rows))

    def test_build_feature_frame_creates_numeric_matrix(self):
        features = build_feature_frame(self.df)

        self.assertEqual(len(self.df), len(features))
        self.assertTrue(all(dtype.kind in "biufc" for dtype in features.dtypes))
        self.assertTrue(any(col.startswith("equipment_id=") for col in features.columns))

    def test_build_supervised_dataset_targets_defect_type(self):
        X, y = build_supervised_dataset(self.df, target_defect_type="气泡")

        self.assertEqual(len(self.df), len(X))
        self.assertEqual(24, int(y.sum()))

    def test_train_random_forest_and_logistic_regression(self):
        X, y = build_supervised_dataset(self.df, target_defect_type="气泡")

        rf = train_tabular_model("random_forest", X, y)
        lr = train_tabular_model("logistic_regression", X, y)

        self.assertIn("model", rf)
        self.assertIn("metrics", rf)
        self.assertIn("model", lr)
        self.assertGreaterEqual(rf["metrics"]["train_accuracy"], 0.5)

    def test_train_isolation_forest_outputs_anomaly_scores(self):
        X = build_feature_frame(self.df)

        result = train_tabular_model("isolation_forest", X)

        self.assertIn("anomaly_scores", result)
        self.assertEqual(len(self.df), len(result["anomaly_scores"]))

    def test_predict_key_factors_returns_model_scores(self):
        predictions = predict_key_factors(self.df, target_defect_type="气泡")

        self.assertFalse(predictions.empty)
        self.assertIn("ml_probability", predictions.columns)
        self.assertIn("model_name", predictions.columns)

    def test_optional_backends_are_reported_without_import_failure(self):
        backends = detect_optional_model_backends()

        self.assertIn("xgboost", backends)
        self.assertIn("lightgbm", backends)

    def test_image_model_wrapper_is_explicitly_unavailable_without_backend(self):
        wrapper = ImageModelWrapper()

        with self.assertRaises(ImageModelUnavailable):
            wrapper.predict([])


if __name__ == "__main__":
    unittest.main()