Spaces:
Running
Running
File size: 3,589 Bytes
7f22d3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import unittest
from unittest.mock import MagicMock, patch
import sys
import os
# Adjust path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from system_manager import SystemManager, SPACE_X, SPACE_R
class TestSystemManager(unittest.TestCase):
def setUp(self):
self.mgr = SystemManager()
@patch('system_manager.client')
@patch('system_manager.get_embedding')
def test_add_to_space_x(self, mock_get_embedding, mock_client):
# Mock embedding
mock_get_embedding.return_value = [0.1]*512
self.mgr.add_to_space_x(text="hello world", url="http://test.com")
mock_client.upsert.assert_called_once()
call_args = mock_client.upsert.call_args
self.assertEqual(call_args.kwargs['collection_name'], SPACE_X)
# Verify payload
points = call_args.kwargs['points']
self.assertEqual(len(points), 1)
self.assertEqual(points[0].payload['url'], "http://test.com")
@patch('system_manager.client')
def test_promote_from_x_to_r(self, mock_client):
# Mock retrieve
mock_point = MagicMock()
mock_point.id = "test_id"
mock_point.vector = [0.1]*512
mock_point.payload = {"url": "http://test.com", "type": "text"}
mock_client.retrieve.return_value = [mock_point]
# Mock search in R (novelty check) - wait, promote_from_x_to_r doesn't check novelty,
# it's an admin force promote.
# It calls trigger_global_recalculation which calls client.scroll
# Mock scroll for trigger_global_recalculation
mock_client.scroll.return_value = ([], None) # Empty R space to stop recursion
success = self.mgr.promote_from_x_to_r("test_id")
self.assertTrue(success)
# Should upsert to SPACE_R
mock_client.upsert.assert_called()
# Check if any upsert call was for SPACE_R
calls = mock_client.upsert.call_args_list
r_calls = [c for c in calls if c.kwargs['collection_name'] == SPACE_R]
self.assertTrue(len(r_calls) > 0)
@patch('system_manager.client')
@patch('system_manager.crawler')
@patch('system_manager.get_embedding')
def test_process_url_and_add(self, mock_get_embedding, mock_crawler, mock_client):
# Mock crawler
mock_crawler.parse.return_value = {
'texts': ["Content 1", "Content 2"],
'images': []
}
# Mock embedding
mock_get_embedding.return_value = [0.1]*512
# Mock novelty check.
# _check_novelty uses self.r_cache.
# We can mock _check_novelty method on the instance since it is a method.
with patch.object(self.mgr, '_check_novelty', return_value=(True, 1.0)):
# Mock trigger_global_recalculation to avoid complex logic
with patch.object(self.mgr, 'trigger_global_recalculation') as mock_trigger:
self.mgr.process_url_and_add("http://test.com")
# Should call upsert to SPACE_R (because novel)
calls = mock_client.upsert.call_args_list
r_calls = [c for c in calls if c.kwargs['collection_name'] == SPACE_R]
self.assertTrue(len(r_calls) > 0)
# Should also call upsert to SPACE_X
x_calls = [c for c in calls if c.kwargs['collection_name'] == SPACE_X]
self.assertTrue(len(x_calls) > 0)
if __name__ == '__main__':
unittest.main()
|