143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
"""
|
|
©AngelaMos | 2026
|
|
test_inference.py
|
|
|
|
Tests the InferenceEngine: model loading, predict output shapes, score ranges, and missing-model handling.
|
|
"""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from ml.autoencoder import ThreatAutoencoder
|
|
from ml.export_onnx import (
|
|
export_autoencoder,
|
|
export_isolation_forest,
|
|
export_random_forest,
|
|
)
|
|
from ml.scaler import FeatureScaler
|
|
from sklearn.ensemble import IsolationForest, RandomForestClassifier
|
|
|
|
from app.core.detection.inference import InferenceEngine
|
|
|
|
|
|
@pytest.fixture
|
|
def model_dir(tmp_path: Path) -> Path:
|
|
"""
|
|
Create a temp directory with all 3 ONNX models + scaler + threshold
|
|
"""
|
|
rng = np.random.default_rng(42)
|
|
X = rng.standard_normal((200, 35)).astype(np.float32)
|
|
y = np.concatenate([np.zeros(140, dtype=int), np.ones(60, dtype=int)])
|
|
|
|
ae = ThreatAutoencoder(input_dim=35)
|
|
export_autoencoder(ae, tmp_path / "ae.onnx")
|
|
|
|
rf = RandomForestClassifier(n_estimators=10, random_state=42)
|
|
rf.fit(X, y)
|
|
export_random_forest(rf, 35, tmp_path / "rf.onnx")
|
|
|
|
iso = IsolationForest(n_estimators=10, random_state=42)
|
|
iso.fit(X[:140])
|
|
export_isolation_forest(iso, 35, tmp_path / "if.onnx")
|
|
|
|
scaler = FeatureScaler()
|
|
scaler.fit(X[:140])
|
|
scaler.save_json(tmp_path / "scaler.json")
|
|
|
|
threshold_data = {"threshold": 0.05}
|
|
(tmp_path / "threshold.json").write_text(json.dumps(threshold_data))
|
|
|
|
return tmp_path
|
|
|
|
|
|
class TestInferenceEngine:
|
|
|
|
def test_loads_all_models(self, model_dir: Path) -> None:
|
|
"""
|
|
Engine reports is_loaded=True when all three ONNX models are present.
|
|
"""
|
|
engine = InferenceEngine(model_dir=str(model_dir))
|
|
assert engine.is_loaded
|
|
|
|
def test_returns_none_when_no_models(self) -> None:
|
|
"""
|
|
Engine reports is_loaded=False when the model directory does not exist.
|
|
"""
|
|
engine = InferenceEngine(model_dir="/nonexistent/path")
|
|
assert not engine.is_loaded
|
|
|
|
def test_predict_returns_none_when_not_loaded(self) -> None:
|
|
"""
|
|
predict returns None when the engine has no models loaded.
|
|
"""
|
|
engine = InferenceEngine(model_dir="/nonexistent/path")
|
|
result = engine.predict(np.zeros((1, 35), dtype=np.float32))
|
|
assert result is None
|
|
|
|
def test_predict_returns_scores(self, model_dir: Path) -> None:
|
|
"""
|
|
predict returns a dict with ae, rf, and if score arrays.
|
|
"""
|
|
engine = InferenceEngine(model_dir=str(model_dir))
|
|
rng = np.random.default_rng(99)
|
|
x = rng.standard_normal((4, 35)).astype(np.float32)
|
|
result = engine.predict(x)
|
|
assert result is not None
|
|
assert "ae" in result
|
|
assert "rf" in result
|
|
assert "if" in result
|
|
|
|
def test_predict_ae_scores_are_positive(self, model_dir: Path) -> None:
|
|
"""
|
|
AE reconstruction error scores are non-negative for all samples.
|
|
"""
|
|
engine = InferenceEngine(model_dir=str(model_dir))
|
|
rng = np.random.default_rng(99)
|
|
x = rng.standard_normal((4, 35)).astype(np.float32)
|
|
result = engine.predict(x)
|
|
assert result is not None
|
|
assert all(s >= 0.0 for s in result["ae"])
|
|
|
|
def test_predict_rf_probabilities_in_range(self, model_dir: Path) -> None:
|
|
"""
|
|
RF malicious-class probabilities are within [0, 1].
|
|
"""
|
|
engine = InferenceEngine(model_dir=str(model_dir))
|
|
rng = np.random.default_rng(99)
|
|
x = rng.standard_normal((4, 35)).astype(np.float32)
|
|
result = engine.predict(x)
|
|
assert result is not None
|
|
assert all(0.0 <= p <= 1.0 for p in result["rf"])
|
|
|
|
def test_predict_single_sample(self, model_dir: Path) -> None:
|
|
"""
|
|
predict works on a single sample and returns one score per model.
|
|
"""
|
|
engine = InferenceEngine(model_dir=str(model_dir))
|
|
rng = np.random.default_rng(99)
|
|
x = rng.standard_normal((1, 35)).astype(np.float32)
|
|
result = engine.predict(x)
|
|
assert result is not None
|
|
assert len(result["ae"]) == 1
|
|
assert len(result["rf"]) == 1
|
|
assert len(result["if"]) == 1
|
|
|
|
def test_threshold_loaded(self, model_dir: Path) -> None:
|
|
"""
|
|
Autoencoder threshold is read from threshold.json on initialization.
|
|
"""
|
|
engine = InferenceEngine(model_dir=str(model_dir))
|
|
assert engine.threshold == 0.05
|
|
|
|
def test_partial_models_not_loaded(self, tmp_path: Path) -> None:
|
|
"""
|
|
Engine with only the AE model present reports is_loaded=False.
|
|
"""
|
|
ae = ThreatAutoencoder(input_dim=35)
|
|
export_autoencoder(ae, tmp_path / "ae.onnx")
|
|
engine = InferenceEngine(model_dir=str(tmp_path))
|
|
assert not engine.is_loaded
|