import unittest from unittest.mock import MagicMock, patch import sys import os # Add project root to path sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) # Mock dependencies to avoid import hangs and external calls sys.modules['spacy'] = MagicMock() sys.modules['semantica.semantic_extract.methods'] = MagicMock() sys.modules['semantica.utils.logging'] = MagicMock() sys.modules['semantica.utils.progress_tracker'] = MagicMock() sys.modules['semantica.semantic_extract.providers'] = MagicMock() # Mock get_logger and get_progress_tracker mock_logger = MagicMock() sys.modules['semantica.utils.logging'].get_logger.return_value = mock_logger mock_tracker = MagicMock() sys.modules['semantica.utils.progress_tracker'].get_progress_tracker.return_value = mock_tracker # Mock the methods module functions specifically mock_methods = sys.modules['semantica.semantic_extract.methods'] mock_methods.get_entity_method = MagicMock() mock_methods.get_relation_method = MagicMock() mock_methods.get_triplet_method = MagicMock() # Mock specific extraction functions mock_extract_entities_hf = MagicMock() mock_extract_relations_hf = MagicMock() mock_extract_triplets_hf = MagicMock() # Setup the registry mocks to return our mock functions mock_methods.get_entity_method.return_value = mock_extract_entities_hf mock_methods.get_relation_method.return_value = mock_extract_relations_hf mock_methods.get_triplet_method.return_value = mock_extract_triplets_hf # Now import the classes under test # We need to patch where they import 'methods' locally if they do with patch.dict(sys.modules): from semantica.semantic_extract.ner_extractor import NERExtractor from semantica.semantic_extract.relation_extractor import RelationExtractor from semantica.semantic_extract.triplet_extractor import TripletExtractor from semantica.semantic_extract.ner_extractor import Entity from semantica.semantic_extract.relation_extractor import Relation class TestExtractorsDispatch(unittest.TestCase): def setUp(self): self.mock_extract_entities_hf = mock_extract_entities_hf self.mock_extract_relations_hf = mock_extract_relations_hf self.mock_extract_triplets_hf = mock_extract_triplets_hf self.mock_extract_entities_hf.reset_mock() self.mock_extract_relations_hf.reset_mock() self.mock_extract_triplets_hf.reset_mock() # Configure mocks to return something iterable/valid self.mock_extract_entities_hf.return_value = [MagicMock(spec=Entity, confidence=0.9, text="Test Entity")] self.mock_extract_relations_hf.return_value = [MagicMock(spec=Relation, confidence=0.9)] self.mock_extract_triplets_hf.return_value = [MagicMock(confidence=0.9)] def test_ner_extractor_huggingface_dispatch(self): print("\nTesting NERExtractor dispatch to HuggingFace...") # Initialize with HuggingFace method extractor = NERExtractor(method="huggingface") # Call extract_entities text = "Steve Jobs founded Apple." # Use a specific model via kwargs extractor.extract_entities(text, model="my-custom-ner-model") # Verify get_entity_method was called with "huggingface" mock_methods.get_entity_method.assert_called_with("huggingface") # Verify the extraction function was called with correct model # We need to check the call args to see if 'model' was passed correctly # The logic we implemented: method_options["model"] = all_options.get("huggingface_model") or all_options.get("model") or self.huggingface_model call_args = self.mock_extract_entities_hf.call_args self.assertIsNotNone(call_args, "extract_entities_huggingface should have been called") _, kwargs = call_args self.assertEqual(kwargs.get("model"), "my-custom-ner-model", "Should use model passed in kwargs") print("NERExtractor dispatch verified.") def test_relation_extractor_huggingface_dispatch(self): print("\nTesting RelationExtractor dispatch to HuggingFace...") extractor = RelationExtractor(method="huggingface") text = "Steve Jobs founded Apple." entities = [MagicMock(spec=Entity)] # Call extract_relations with explicit model extractor.extract_relations(text, entities, model="my-relation-model") # Verify dispatch mock_methods.get_relation_method.assert_called_with("huggingface") call_args = self.mock_extract_relations_hf.call_args self.assertIsNotNone(call_args, "extract_relations_huggingface should have been called") _, kwargs = call_args self.assertEqual(kwargs.get("model"), "my-relation-model", "Should use model passed in kwargs") print("RelationExtractor dispatch verified.") def test_triplet_extractor_huggingface_dispatch(self): print("\nTesting TripletExtractor dispatch to HuggingFace...") extractor = TripletExtractor(method="huggingface") text = "Steve Jobs founded Apple." # Call extract_triplets with explicit model extractor.extract_triplets(text, model="my-triplet-model") # Verify dispatch mock_methods.get_triplet_method.assert_called_with("huggingface") call_args = self.mock_extract_triplets_hf.call_args self.assertIsNotNone(call_args, "extract_triplets_huggingface should have been called") _, kwargs = call_args self.assertEqual(kwargs.get("model"), "my-triplet-model", "Should use model passed in kwargs") print("TripletExtractor dispatch verified.") def test_ner_extractor_huggingface_fallback(self): print("\nTesting NERExtractor fallback logic...") # Init with huggingface_model in config extractor = NERExtractor(method="huggingface", huggingface_model="config-model") extractor.extract_entities("text") _, kwargs = self.mock_extract_entities_hf.call_args self.assertEqual(kwargs.get("model"), "config-model", "Should prioritize huggingface_model from config") # Now override with kwargs model extractor.extract_entities("text", model="kwargs-model") _, kwargs = self.mock_extract_entities_hf.call_args self.assertEqual(kwargs.get("model"), "kwargs-model", "Should allow overriding config huggingface_model via model kwarg") # Let's test passing 'huggingface_model' in kwargs extractor.extract_entities("text", huggingface_model="override-model") _, kwargs = self.mock_extract_entities_hf.call_args self.assertEqual(kwargs.get("model"), "override-model", "Should allow overriding huggingface_model via kwargs") def test_triplet_extractor_lazy_loading(self): print("\nTesting TripletExtractor lazy loading for HuggingFace...") # Initialize with HuggingFace method extractor = TripletExtractor(method="huggingface") # Check initial state self.assertIsNone(extractor._ner_extractor) self.assertIsNone(extractor._relation_extractor) # Run extraction extractor.extract_triplets("Steve Jobs founded Apple.") # Check state AFTER extraction - should STILL be None because huggingface (REBEL) doesn't need them self.assertIsNone(extractor._ner_extractor, "NERExtractor should not be initialized for HuggingFace method") self.assertIsNone(extractor._relation_extractor, "RelationExtractor should not be initialized for HuggingFace method") print("TripletExtractor lazy loading verified.") if __name__ == '__main__': unittest.main()