Spaces:
Sleeping
Sleeping
| import unittest | |
| from unittest.mock import patch, MagicMock, Mock | |
| import sys | |
| import os | |
| import torch | |
| from PIL import Image | |
| # Cấu hình đường dẫn | |
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) | |
| import shared_resources | |
| class TestSharedResources(unittest.TestCase): | |
| def setUp(self): | |
| print("\n" + "-"*60) | |
| shared_resources.db = None | |
| shared_resources.client = None | |
| shared_resources.model = None | |
| shared_resources.processor = None | |
| # ================================================================= | |
| # 1. TEST CLASS MODEL WRAPPER | |
| # ================================================================= | |
| def test_01_fine_tuned_dinov2_init(self, mock_dinov2): | |
| """Test khởi tạo class FineTunedDINOv2 wrapper.""" | |
| print("➤ [TEST 01] Kiểm tra khởi tạo FineTunedDINOv2...") | |
| mock_backbone = Mock() | |
| mock_dinov2.from_pretrained.return_value = mock_backbone | |
| model_instance = shared_resources.FineTunedDINOv2() | |
| self.assertIsInstance(model_instance, torch.nn.Module) | |
| print(" ✔ PASSED: Class khởi tạo thành công.") | |
| # ================================================================= | |
| # 2. TEST LOAD RESOURCES | |
| # ================================================================= | |
| def test_02_load_resources_success(self, mock_processor, mock_cls, mock_exists, mock_torch_load, mock_mongo): | |
| """Test load_resources chạy thành công (Happy Path).""" | |
| print("➤ [TEST 02] Kiểm tra load_resources (Success Case)...") | |
| # --- [SỬA QUAN TRỌNG] Dùng MagicMock để hỗ trợ cú pháp client['db_name'] --- | |
| mock_client_instance = MagicMock() | |
| mock_db_instance = MagicMock() | |
| # Cấu hình: MongoClient() -> trả về client giả | |
| mock_mongo.return_value = mock_client_instance | |
| # Cấu hình: client['tên_db'] -> trả về db giả | |
| mock_client_instance.__getitem__.return_value = mock_db_instance | |
| # Mock file model | |
| mock_exists.return_value = True | |
| mock_checkpoint = {"model_state_dict": {"module.layer1": 1}} | |
| mock_torch_load.return_value = mock_checkpoint | |
| mock_model_instance = Mock() | |
| mock_cls.return_value = mock_model_instance | |
| # --- ACTION --- | |
| shared_resources.load_resources() | |
| # --- ASSERT --- | |
| self.assertIsNotNone(shared_resources.db) | |
| self.assertIsNotNone(shared_resources.client) | |
| # Kiểm tra lệnh tạo index (truy cập qua chuỗi magic mock) | |
| # db['collection'].create_index(...) | |
| mock_db_instance.__getitem__.return_value.create_index.assert_called() | |
| self.assertIsNotNone(shared_resources.model) | |
| print(" ✔ PASSED: DB connected & Model loaded.") | |
| def test_03_load_resources_db_fail(self, mock_mongo): | |
| """Test load_resources khi kết nối DB thất bại.""" | |
| print("➤ [TEST 03] Kiểm tra load_resources (DB Fail)...") | |
| mock_mongo.side_effect = Exception("Connection Timeout") | |
| shared_resources.load_resources() | |
| self.assertIsNone(shared_resources.db) | |
| print(" ✔ PASSED: App handled DB error gracefully.") | |
| def test_04_load_resources_model_fail(self, mock_cls, mock_mongo): | |
| """Test load_resources khi load Model thất bại.""" | |
| print("➤ [TEST 04] Kiểm tra load_resources (Model Fail)...") | |
| # --- [SỬA QUAN TRỌNG] Dùng MagicMock --- | |
| # Nếu dùng Mock thường, dòng db = client[DB_NAME] sẽ lỗi TypeError | |
| mock_mongo.return_value = MagicMock() | |
| # Mock Model thất bại | |
| mock_cls.side_effect = Exception("CUDA Out of memory") | |
| shared_resources.load_resources() | |
| self.assertIsNotNone(shared_resources.db) # DB phải OK | |
| self.assertIsNone(shared_resources.model) # Model phải None | |
| print(" ✔ PASSED: DB OK, Model Failed (Correct behavior).") | |
| # ================================================================= | |
| # 3. TEST EMBED IMAGE | |
| # ================================================================= | |
| def test_05_embed_image_success(self): | |
| """Test hàm embed_image tạo ra vector.""" | |
| print("➤ [TEST 05] Kiểm tra embed_image (Success)...") | |
| mock_model = Mock() | |
| mock_processor = Mock() | |
| shared_resources.model = mock_model | |
| shared_resources.processor = mock_processor | |
| shared_resources.DEVICE = "cpu" | |
| mock_output = torch.rand(1, 768) | |
| mock_model.return_value = mock_output | |
| mock_processor_output = Mock() | |
| mock_processor_output.pixel_values = torch.zeros(1, 3, 224, 224) | |
| mock_processor.return_value = mock_processor_output | |
| result = shared_resources.embed_image(Image.new('RGB', (100, 100))) | |
| self.assertIsInstance(result, list) | |
| self.assertEqual(len(result), 768) | |
| print(" ✔ PASSED: Embed OK.") | |
| def test_06_embed_image_no_model(self): | |
| """Test embed_image báo lỗi khi model chưa load.""" | |
| print("➤ [TEST 06] Kiểm tra embed_image (Model not loaded)...") | |
| shared_resources.model = None | |
| with self.assertRaises(RuntimeError): | |
| shared_resources.embed_image(Image.new('RGB', (10, 10))) | |
| print(" ✔ PASSED: RuntimeError raised.") | |
| # ================================================================= | |
| # 4. TEST VECTOR SEARCH | |
| # ================================================================= | |
| def test_07_search_similar_landmark_success(self): | |
| """Test hàm search gọi MongoDB aggregation đúng cách.""" | |
| print("➤ [TEST 07] Kiểm tra search_similar_landmark (Aggregation)...") | |
| # --- [SỬA QUAN TRỌNG] Dùng MagicMock cho DB --- | |
| mock_db = MagicMock() | |
| mock_col = MagicMock() | |
| # Cấu hình: db['collection_name'] -> trả về collection giả | |
| mock_db.__getitem__.return_value = mock_col | |
| shared_resources.db = mock_db | |
| mock_results = [{"landmark_id": "123", "score": 0.9}] | |
| mock_col.aggregate.return_value = mock_results | |
| dummy_vector = [0.1] * 768 | |
| result = shared_resources.search_similar_landmark(dummy_vector, "my_collection") | |
| self.assertEqual(result, mock_results) | |
| mock_col.aggregate.assert_called_once() | |
| print(" ✔ PASSED: Aggregation called correctly.") | |
| def test_08_search_no_db(self): | |
| """Test search báo lỗi khi DB chưa kết nối.""" | |
| print("➤ [TEST 08] Kiểm tra search_similar_landmark (No DB)...") | |
| shared_resources.db = None | |
| with self.assertRaises(RuntimeError): | |
| shared_resources.search_similar_landmark([], "col") | |
| print(" ✔ PASSED: RuntimeError raised.") | |
| if __name__ == '__main__': | |
| unittest.main(verbosity=2) |