Add unit tests for PhishingDetectionPipeline class

tests/test_inference_pipeline.py

@@ -0,0 +1,441 @@
+"""Tests for inference pipeline."""
+
+import pytest
+import tempfile
+import os
+from unittest.mock import Mock, patch, MagicMock
+import pandas as pd
+import numpy as np
+
+from src.phising_detection.inference.pipeline import PhishingDetectionPipeline
+
+
+class TestPhishingDetectionPipeline:
+    """Tests for PhishingDetectionPipeline class."""
+
+    def test_init_without_urlscan(self):
+        """Test initialization without URLScan API key."""
+        pipeline = PhishingDetectionPipeline(
+            model_name="test_model",
+            model_version=1
+        )
+        assert pipeline.model_name == "test_model"
+        assert pipeline.model_version == 1
+        assert pipeline.model is None
+        assert pipeline.scaler is None
+        assert pipeline.feature_names is None
+        assert pipeline.urlscan_client is None
+
+    def test_init_with_urlscan(self):
+        """Test initialization with URLScan API key."""
+        with patch('src.phising_detection.inference.pipeline.URLScanClient') as mock_client:
+            pipeline = PhishingDetectionPipeline(
+                model_name="test_model",
+                urlscan_api_key="test_key"
+            )
+            mock_client.assert_called_once_with(api_key="test_key")
+            assert pipeline.urlscan_client is not None
+
+    def test_is_loaded_false(self):
+        """Test is_loaded returns False when model not loaded."""
+        pipeline = PhishingDetectionPipeline()
+        assert pipeline.is_loaded() is False
+
+    def test_is_loaded_true(self):
+        """Test is_loaded returns True when model is loaded."""
+        pipeline = PhishingDetectionPipeline()
+        pipeline.model = Mock()
+        pipeline.scaler = Mock()
+        pipeline.feature_names = ['feature1', 'feature2']
+        assert pipeline.is_loaded() is True
+
+    @patch('src.phising_detection.inference.pipeline.connect_to_hopsworks')
+    @patch('src.phising_detection.inference.pipeline.joblib.load')
+    def test_load_model_from_hopsworks(self, mock_joblib_load, mock_connect):
+        """Test loading model from Hopsworks."""
+        # Setup mocks
+        mock_project = Mock()
+        mock_mr = Mock()
+        mock_model_registry = Mock()
+        mock_model_registry.version = 1
+        mock_model_registry.download.return_value = "/tmp/model_dir"
+
+        mock_project.get_model_registry.return_value = mock_mr
+        mock_mr.get_model.return_value = mock_model_registry
+        mock_connect.return_value = mock_project
+
+        # Mock model and scaler
+        mock_model = Mock()
+        mock_scaler = Mock()
+        mock_joblib_load.side_effect = [mock_model, mock_scaler]
+
+        # Create temporary feature_names file
+        with tempfile.TemporaryDirectory() as tmpdir:
+            feature_names_path = os.path.join(tmpdir, "feature_names.txt")
+            with open(feature_names_path, 'w') as f:
+                f.write("feature1\nfeature2\nfeature3")
+
+            # Mock download to return our temp dir
+            mock_model_registry.download.return_value = tmpdir
+
+            # Test
+            pipeline = PhishingDetectionPipeline(model_name="test_model")
+            pipeline.load_model_from_hopsworks()
+
+            # Assertions
+            assert pipeline.model == mock_model
+            assert pipeline.scaler == mock_scaler
+            assert pipeline.feature_names == ['feature1', 'feature2', 'feature3']
+            mock_connect.assert_called_once()
+            mock_mr.get_model.assert_called_once_with("test_model")
+
+    @patch('src.phising_detection.inference.pipeline.connect_to_hopsworks')
+    @patch('src.phising_detection.inference.pipeline.joblib.load')
+    def test_load_model_with_version(self, mock_joblib_load, mock_connect):
+        """Test loading specific model version from Hopsworks."""
+        # Setup mocks
+        mock_project = Mock()
+        mock_mr = Mock()
+        mock_model_registry = Mock()
+        mock_model_registry.version = 2
+        mock_model_registry.download.return_value = "/tmp/model_dir"
+
+        mock_project.get_model_registry.return_value = mock_mr
+        mock_mr.get_model.return_value = mock_model_registry
+        mock_connect.return_value = mock_project
+
+        mock_joblib_load.side_effect = [Mock(), Mock()]
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            feature_names_path = os.path.join(tmpdir, "feature_names.txt")
+            with open(feature_names_path, 'w') as f:
+                f.write("feature1")
+
+            mock_model_registry.download.return_value = tmpdir
+
+            # Test with specific version
+            pipeline = PhishingDetectionPipeline(model_name="test_model", model_version=2)
+            pipeline.load_model_from_hopsworks()
+
+            # Should request version 2
+            mock_mr.get_model.assert_called_once_with("test_model", version=2)
+
+    def test_preprocess_features_not_loaded(self):
+        """Test preprocessing fails when model not loaded."""
+        pipeline = PhishingDetectionPipeline()
+        features = {'feature1': 10, 'feature2': 20}
+
+        with pytest.raises(ValueError) as exc_info:
+            pipeline.preprocess_features(features)
+        assert "Model not loaded" in str(exc_info.value)
+
+    def test_preprocess_features_success(self):
+        """Test successful feature preprocessing."""
+        # Setup pipeline with mock components
+        pipeline = PhishingDetectionPipeline()
+        pipeline.model = Mock()
+        pipeline.feature_names = [
+            'domain_age_days',
+            'secure_percentage',
+            'has_umbrella_rank',
+            'umbrella_rank',
+            'has_tls',
+            'tls_valid_days',
+            'url_length',
+            'subdomain_count'
+        ]
+
+        # Mock scaler
+        mock_scaler = Mock()
+        mock_scaler.transform.return_value = np.array([[1.5, 0.8, 5000, 365, 25, 1]])
+        pipeline.scaler = mock_scaler
+
+        # Test features
+        features = {
+            'domain_age_days': 3000,
+            'secure_percentage': 95.0,
+            'has_umbrella_rank': 1,
+            'umbrella_rank': 5000,
+            'has_tls': 1,
+            'tls_valid_days': 365,
+            'url_length': 25,
+            'subdomain_count': 1
+        }
+
+        result = pipeline.preprocess_features(features)
+
+        # Assertions
+        assert isinstance(result, pd.DataFrame)
+        assert list(result.columns) == pipeline.feature_names
+        assert len(result) == 1
+        mock_scaler.transform.assert_called_once()
+
+    def test_preprocess_features_with_missing_features(self):
+        """Test preprocessing handles missing features."""
+        pipeline = PhishingDetectionPipeline()
+        pipeline.model = Mock()
+        pipeline.feature_names = [
+            'domain_age_days', 'secure_percentage', 'has_umbrella_rank',
+            'umbrella_rank', 'has_tls', 'tls_valid_days', 'url_length', 'subdomain_count'
+        ]
+
+        mock_scaler = Mock()
+        # Mock the scaler to return the same shape as input continuous features (6 features)
+        mock_scaler.transform.return_value = np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]])
+        pipeline.scaler = mock_scaler
+
+        # Only provide some features (missing subdomain_count)
+        features = {
+            'domain_age_days': 3000,
+            'secure_percentage': 95.0,
+            'has_umbrella_rank': 1,
+            'umbrella_rank': 5000,
+            'has_tls': 1,
+            'tls_valid_days': 365,
+            'url_length': 25
+        }
+
+        result = pipeline.preprocess_features(features)
+
+        # Should add missing subdomain_count and all features should be present
+        assert 'subdomain_count' in result.columns
+        assert len(result.columns) == 8  # All 8 features should be present
+        assert list(result.columns) == pipeline.feature_names
+
+    def test_predict_not_loaded(self):
+        """Test prediction fails when model not loaded."""
+        pipeline = PhishingDetectionPipeline()
+        features = {'feature1': 10}
+
+        with pytest.raises(ValueError) as exc_info:
+            pipeline.predict(features)
+        assert "Model not loaded" in str(exc_info.value)
+
+    def test_predict_phishing(self):
+        """Test prediction for phishing URL."""
+        # Setup pipeline
+        pipeline = PhishingDetectionPipeline()
+        pipeline.feature_names = [
+            'domain_age_days', 'secure_percentage', 'has_umbrella_rank',
+            'umbrella_rank', 'has_tls', 'tls_valid_days', 'url_length', 'subdomain_count'
+        ]
+
+        # Mock model
+        mock_model = Mock()
+        mock_model.predict_proba.return_value = np.array([[0.2, 0.8]])  # 80% phishing
+        mock_model.predict.return_value = np.array([1])  # Phishing
+        pipeline.model = mock_model
+
+        # Mock scaler
+        mock_scaler = Mock()
+        mock_scaler.transform.return_value = np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]])
+        pipeline.scaler = mock_scaler
+
+        # Test features
+        features = {
+            'domain_age_days': 10,
+            'secure_percentage': 50.0,
+            'has_umbrella_rank': 0,
+            'umbrella_rank': 999999,
+            'has_tls': 0,
+            'tls_valid_days': 0,
+            'url_length': 150,
+            'subdomain_count': 5
+        }
+
+        result = pipeline.predict(features)
+
+        # Assertions
+        assert result['prediction'] == "PHISHING"
+        assert result['is_phishing'] is True
+        assert result['confidence'] == 0.8
+        assert result['phishing_probability'] == 0.8
+        assert result['legitimate_probability'] == 0.2
+
+    def test_predict_legitimate(self):
+        """Test prediction for legitimate URL."""
+        # Setup pipeline
+        pipeline = PhishingDetectionPipeline()
+        pipeline.feature_names = [
+            'domain_age_days', 'secure_percentage', 'has_umbrella_rank',
+            'umbrella_rank', 'has_tls', 'tls_valid_days', 'url_length', 'subdomain_count'
+        ]
+
+        # Mock model
+        mock_model = Mock()
+        mock_model.predict_proba.return_value = np.array([[0.9, 0.1]])  # 90% legitimate
+        mock_model.predict.return_value = np.array([0])  # Legitimate
+        pipeline.model = mock_model
+
+        # Mock scaler
+        mock_scaler = Mock()
+        mock_scaler.transform.return_value = np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]])
+        pipeline.scaler = mock_scaler
+
+        # Test features
+        features = {
+            'domain_age_days': 3000,
+            'secure_percentage': 95.0,
+            'has_umbrella_rank': 1,
+            'umbrella_rank': 5000,
+            'has_tls': 1,
+            'tls_valid_days': 365,
+            'url_length': 25,
+            'subdomain_count': 1
+        }
+
+        result = pipeline.predict(features)
+
+        # Assertions
+        assert result['prediction'] == "LEGITIMATE"
+        assert result['is_phishing'] is False
+        assert result['confidence'] == 0.9
+        assert result['phishing_probability'] == 0.1
+        assert result['legitimate_probability'] == 0.9
+
+    def test_predict_url_without_urlscan(self):
+        """Test predict_url fails without URLScan client."""
+        pipeline = PhishingDetectionPipeline()
+        pipeline.model = Mock()
+        pipeline.scaler = Mock()
+        pipeline.feature_names = ['feature1']
+
+        with pytest.raises(ValueError) as exc_info:
+            pipeline.predict_url("https://example.com")
+        assert "URLScan client not initialized" in str(exc_info.value)
+
+    @patch('src.phising_detection.inference.pipeline.extract_features')
+    def test_predict_url_success(self, mock_extract_features):
+        """Test successful end-to-end URL prediction."""
+        # Setup pipeline
+        pipeline = PhishingDetectionPipeline()
+        pipeline.feature_names = [
+            'domain_age_days', 'secure_percentage', 'has_umbrella_rank',
+            'umbrella_rank', 'has_tls', 'tls_valid_days', 'url_length', 'subdomain_count'
+        ]
+
+        # Mock model
+        mock_model = Mock()
+        mock_model.predict_proba.return_value = np.array([[0.7, 0.3]])
+        mock_model.predict.return_value = np.array([0])
+        pipeline.model = mock_model
+
+        # Mock scaler
+        mock_scaler = Mock()
+        mock_scaler.transform.return_value = np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]])
+        pipeline.scaler = mock_scaler
+
+        # Mock URLScan client
+        mock_urlscan_client = Mock()
+        mock_scan_result = {
+            'task': {'uuid': 'test-uuid-123'},
+            'page': {'domainAgeDays': 3000},
+            'stats': {'securePercentage': 95}
+        }
+        mock_urlscan_client.submit_and_wait.return_value = mock_scan_result
+        pipeline.urlscan_client = mock_urlscan_client
+
+        # Mock extracted features
+        extracted_features = {
+            'domain_age_days': 3000,
+            'secure_percentage': 95.0,
+            'has_umbrella_rank': 1,
+            'umbrella_rank': 5000,
+            'has_tls': 1,
+            'tls_valid_days': 365,
+            'url_length': 25,
+            'subdomain_count': 1
+        }
+        mock_extract_features.return_value = extracted_features
+
+        # Test
+        result = pipeline.predict_url("https://example.com")
+
+        # Assertions
+        assert result['prediction'] == "LEGITIMATE"
+        assert result['confidence'] == 0.7
+        assert result['features'] == extracted_features
+        assert result['scan_uuid'] == 'test-uuid-123'
+        mock_urlscan_client.submit_and_wait.assert_called_once_with("https://example.com")
+
+    def test_predict_url_scan_fails(self):
+        """Test predict_url handles URLScan failure."""
+        # Setup pipeline
+        pipeline = PhishingDetectionPipeline()
+        pipeline.model = Mock()
+        pipeline.scaler = Mock()
+        pipeline.feature_names = ['feature1']
+
+        # Mock URLScan client that returns None
+        mock_urlscan_client = Mock()
+        mock_urlscan_client.submit_and_wait.return_value = None
+        pipeline.urlscan_client = mock_urlscan_client
+
+        # Test
+        result = pipeline.predict_url("https://example.com")
+
+        # Should return error
+        assert 'error' in result
+        assert result['prediction'] == "ERROR"
+        assert result['confidence'] == 0.0
+
+    @patch('src.phising_detection.inference.pipeline.extract_features')
+    def test_predict_url_exception_handling(self, mock_extract_features):
+        """Test predict_url handles exceptions gracefully."""
+        # Setup pipeline
+        pipeline = PhishingDetectionPipeline()
+        pipeline.model = Mock()
+        pipeline.scaler = Mock()
+        pipeline.feature_names = ['feature1']
+
+        # Mock URLScan client
+        mock_urlscan_client = Mock()
+        mock_urlscan_client.submit_and_wait.return_value = {'task': {}}
+        pipeline.urlscan_client = mock_urlscan_client
+
+        # Mock extract_features to raise exception
+        mock_extract_features.side_effect = Exception("Test error")
+
+        # Test
+        result = pipeline.predict_url("https://example.com")
+
+        # Should return error
+        assert 'error' in result
+        assert result['prediction'] == "ERROR"
+        assert "Test error" in result['error']
+
+    def test_predict_numerical_stability(self):
+        """Test prediction handles edge cases in probabilities."""
+        # Setup pipeline
+        pipeline = PhishingDetectionPipeline()
+        pipeline.feature_names = [
+            'domain_age_days', 'secure_percentage', 'has_umbrella_rank',
+            'umbrella_rank', 'has_tls', 'tls_valid_days', 'url_length', 'subdomain_count'
+        ]
+
+        # Mock model with extreme probabilities
+        mock_model = Mock()
+        mock_model.predict_proba.return_value = np.array([[0.999999, 0.000001]])
+        mock_model.predict.return_value = np.array([0])
+        pipeline.model = mock_model
+
+        mock_scaler = Mock()
+        mock_scaler.transform.return_value = np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]])
+        pipeline.scaler = mock_scaler
+
+        features = {
+            'domain_age_days': 3000,
+            'secure_percentage': 95.0,
+            'has_umbrella_rank': 1,
+            'umbrella_rank': 5000,
+            'has_tls': 1,
+            'tls_valid_days': 365,
+            'url_length': 25,
+            'subdomain_count': 1
+        }
+        result = pipeline.predict(features)
+
+        # Should handle extreme values correctly
+        assert result['confidence'] > 0.99
+        assert result['prediction'] == "LEGITIMATE"
+        assert result['phishing_probability'] < 0.01