tests/test_search_engine.py

import unittest
from unittest.mock import MagicMock, patch
import numpy as np
import sys
import os

# Adjust path to import modules from parent directory
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from search_engine import gauss_rank_norm, search

class TestSearchEngine(unittest.TestCase):

    def test_gauss_rank_norm(self):
        scores = [0.1, 0.5, 0.9]
        norm_scores = gauss_rank_norm(scores)
        # Rank data: 1, 2, 3 -> divided by 3 -> 0.33, 0.66, 1.0
        self.assertAlmostEqual(norm_scores[0], 1/3)
        self.assertAlmostEqual(norm_scores[1], 2/3)
        self.assertAlmostEqual(norm_scores[2], 1.0)
        
        # Test empty
        self.assertEqual(gauss_rank_norm([]), [])

    @patch('search_engine.client')
    @patch('search_engine.clip_model')
    @patch('search_engine.clip_processor')
    @patch('search_engine.consistency_engine')
    def test_search(self, mock_consistency, mock_processor, mock_model, mock_client):
        # Mock CLIP
        mock_processor.return_value = {'input_ids': 'fake'}
        
        mock_features = MagicMock()
        # Return a tensor that supports .norm() and .numpy()
        mock_tensor = MagicMock()
        mock_tensor.norm.return_value = 1.0
        mock_tensor.__truediv__.return_value = mock_tensor
        mock_tensor.__getitem__.return_value = mock_tensor
        mock_tensor.numpy.return_value = np.array([[0.1, 0.2]])
        mock_tensor.tolist.return_value = [0.1, 0.2]
        
        mock_model.get_text_features.return_value = mock_tensor

        # Mock Qdrant Search
        mock_hit = MagicMock()
        mock_hit.id = "test_id_1"
        mock_hit.score = 0.9
        mock_hit.payload = {
            "pr_score": 0.5,
            "type": "text",
            "url": "http://example.com",
            "content_preview": "Test Content"
        }
        mock_client.search.return_value = [mock_hit]

        # Mock Consistency Engine
        mock_consistency.check_consistency.return_value = (True, 0.1)

        # Run Search
        results = search("test query", top_k=5)

        # Assertions
        self.assertTrue(len(results) > 0)
        self.assertEqual(results[0]['id'], "test_id_1")
        self.assertEqual(results[0]['url'], "http://example.com")
        
        # Verify calls
        mock_client.search.assert_called_once()
        mock_consistency.check_consistency.assert_called_once()

    @patch('search_engine.client')
    @patch('search_engine.clip_model')
    @patch('search_engine.clip_processor')
    @patch('search_engine.consistency_engine')
    def test_search_consistency_failure(self, mock_consistency, mock_processor, mock_model, mock_client):
        # Mock CLIP (same as above)
        mock_processor.return_value = {'input_ids': 'fake'}
        mock_tensor = MagicMock()
        mock_tensor.norm.return_value = 1.0
        mock_tensor.__truediv__.return_value = mock_tensor
        mock_tensor.__getitem__.return_value = mock_tensor
        mock_tensor.numpy.return_value = np.array([[0.1, 0.2]])
        mock_tensor.tolist.return_value = [0.1, 0.2]
        mock_model.get_text_features.return_value = mock_tensor

        # Mock Qdrant Search
        mock_hit = MagicMock()
        mock_hit.id = "test_id_blocked"
        mock_hit.score = 0.9
        mock_hit.payload = {}
        mock_client.search.return_value = [mock_hit]

        # Mock Consistency Engine to FAIL
        mock_consistency.check_consistency.return_value = (False, 10.0)

        # Run Search
        results = search("test query", top_k=5)

        # Should be empty because the only result was blocked
        self.assertEqual(len(results), 0)

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