From d365113440d6342e68e19c6f38e1502af39c55f3 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Mon, 2 Dec 2024 23:49:30 +0200
Subject: [PATCH] feat(add more tests):

---
 .../refusal_classifier/test_model.py          | 53 +++++++++++++++++++
 agentic_security/routes/test_report.py        |  6 ++-
 agentic_security/routes/test_static.py        | 10 ++--
 3 files changed, 63 insertions(+), 6 deletions(-)
 create mode 100644 agentic_security/refusal_classifier/test_model.py

diff --git a/agentic_security/refusal_classifier/test_model.py b/agentic_security/refusal_classifier/test_model.py
new file mode 100644
index 0000000..b1f017b
--- /dev/null
+++ b/agentic_security/refusal_classifier/test_model.py
@@ -0,0 +1,53 @@
+from unittest.mock import patch
+
+import pandas as pd
+import pytest
+
+from .model import RefusalClassifier
+
+
+@pytest.fixture
+def mock_training_data():
+    """Create mock training data CSV content"""
+    data = {
+        "GPT4_response": ["I cannot help with that", "I must decline"],
+        "ChatGPT_response": ["I won't assist with that", "That's not appropriate"],
+        "Claude_response": ["I cannot comply", "That would be unethical"],
+    }
+    return pd.DataFrame(data)
+
+
+@pytest.fixture
+def classifier():
+    """Create a RefusalClassifier instance with test paths"""
+    return RefusalClassifier(
+        model_path="test_model.joblib",
+        vectorizer_path="test_vectorizer.joblib",
+        scaler_path="test_scaler.joblib",
+    )
+
+
+@pytest.fixture
+def trained_classifier(classifier, mock_training_data):
+    """Create a trained classifier with mock data"""
+    with patch("pandas.read_csv", return_value=mock_training_data):
+        classifier.train(["mock_data.csv"])
+    return classifier
+
+
+def test_is_refusal_without_loading():
+    """Test prediction without loading model raises error"""
+    classifier = RefusalClassifier()
+    with pytest.raises(ValueError, match="Model, vectorizer, or scaler not loaded"):
+        classifier.is_refusal("test text")
+
+
+def test_is_refusal(trained_classifier):
+    """Test refusal prediction"""
+    # Test refusal text
+    refusal_text = "I cannot help with that kind of request"
+    assert trained_classifier.is_refusal(refusal_text) in [True, False]
+
+    # Test non-refusal text
+    normal_text = "Here's the information you requested"
+    assert trained_classifier.is_refusal(normal_text) in [True, False]
diff --git a/agentic_security/routes/test_report.py b/agentic_security/routes/test_report.py
index 7166376..237f262 100644
--- a/agentic_security/routes/test_report.py
+++ b/agentic_security/routes/test_report.py
@@ -1,7 +1,9 @@
+from pathlib import Path
+from unittest.mock import patch
+
 import pytest
 from fastapi.testclient import TestClient
-from unittest.mock import patch, Mock
-from pathlib import Path
+
 from .report import router
 
 client = TestClient(router)
diff --git a/agentic_security/routes/test_static.py b/agentic_security/routes/test_static.py
index 91700f3..fc6ca90 100644
--- a/agentic_security/routes/test_static.py
+++ b/agentic_security/routes/test_static.py
@@ -1,9 +1,11 @@
-import pytest
-from fastapi.testclient import TestClient
 from pathlib import Path
-from ..models.schemas import Settings
-from .static import router, get_static_file
+
+import pytest
 from fastapi import HTTPException
+from fastapi.testclient import TestClient
+
+from ..models.schemas import Settings
+from .static import get_static_file, router
 
 client = TestClient(router)