Fix tests

This commit is contained in:
Cyber MacGeddon 2026-03-13 11:33:50 +00:00
parent e8407b3441
commit fc54bede37
3 changed files with 18 additions and 18 deletions

View file

@ -128,7 +128,7 @@ class TestAgentKgExtractionIntegration:
# Parse and process # Parse and process
extraction_data = extractor.parse_jsonl(agent_response) extraction_data = extractor.parse_jsonl(agent_response)
triples, entity_contexts = extractor.process_extraction_data(extraction_data, v.metadata) triples, entity_contexts, extracted_triples = extractor.process_extraction_data(extraction_data, v.metadata)
# Emit outputs # Emit outputs
if triples: if triples:

View file

@ -168,7 +168,7 @@ This is not JSON at all
} }
] ]
triples, entity_contexts = agent_extractor.process_extraction_data(data, sample_metadata) triples, entity_contexts, _ = agent_extractor.process_extraction_data(data, sample_metadata)
# Check entity label triple # Check entity label triple
label_triple = next((t for t in triples if t.p.iri == RDF_LABEL and t.o.value == "Machine Learning"), None) label_triple = next((t for t in triples if t.p.iri == RDF_LABEL and t.o.value == "Machine Learning"), None)
@ -206,7 +206,7 @@ This is not JSON at all
} }
] ]
triples, entity_contexts = agent_extractor.process_extraction_data(data, sample_metadata) triples, entity_contexts, _ = agent_extractor.process_extraction_data(data, sample_metadata)
# Check that subject, predicate, and object labels are created # Check that subject, predicate, and object labels are created
subject_uri = f"{TRUSTGRAPH_ENTITIES}Machine%20Learning" subject_uri = f"{TRUSTGRAPH_ENTITIES}Machine%20Learning"
@ -244,7 +244,7 @@ This is not JSON at all
} }
] ]
triples, entity_contexts = agent_extractor.process_extraction_data(data, sample_metadata) triples, entity_contexts, _ = agent_extractor.process_extraction_data(data, sample_metadata)
# Check that object labels are not created for literal objects # Check that object labels are not created for literal objects
object_labels = [t for t in triples if t.p.iri == RDF_LABEL and t.o.value == "95%"] object_labels = [t for t in triples if t.p.iri == RDF_LABEL and t.o.value == "95%"]
@ -253,7 +253,7 @@ This is not JSON at all
def test_process_extraction_data_combined(self, agent_extractor, sample_metadata, sample_extraction_data): def test_process_extraction_data_combined(self, agent_extractor, sample_metadata, sample_extraction_data):
"""Test processing of combined definitions and relationships""" """Test processing of combined definitions and relationships"""
triples, entity_contexts = agent_extractor.process_extraction_data(sample_extraction_data, sample_metadata) triples, entity_contexts, _ = agent_extractor.process_extraction_data(sample_extraction_data, sample_metadata)
# Check that we have both definition and relationship triples # Check that we have both definition and relationship triples
definition_triples = [t for t in triples if t.p.iri == DEFINITION] definition_triples = [t for t in triples if t.p.iri == DEFINITION]
@ -272,7 +272,7 @@ This is not JSON at all
{"type": "definition", "entity": "Test Entity", "definition": "Test definition"} {"type": "definition", "entity": "Test Entity", "definition": "Test definition"}
] ]
triples, entity_contexts = agent_extractor.process_extraction_data(data, metadata) triples, entity_contexts, _ = agent_extractor.process_extraction_data(data, metadata)
# Should not create subject-of relationships when no metadata ID # Should not create subject-of relationships when no metadata ID
subject_of_triples = [t for t in triples if t.p.iri == SUBJECT_OF] subject_of_triples = [t for t in triples if t.p.iri == SUBJECT_OF]
@ -285,7 +285,7 @@ This is not JSON at all
"""Test processing of empty extraction data""" """Test processing of empty extraction data"""
data = [] data = []
triples, entity_contexts = agent_extractor.process_extraction_data(data, sample_metadata) triples, entity_contexts, _ = agent_extractor.process_extraction_data(data, sample_metadata)
# Should have no entity contexts # Should have no entity contexts
assert len(entity_contexts) == 0 assert len(entity_contexts) == 0
@ -300,7 +300,7 @@ This is not JSON at all
{"type": "relationship", "subject": "A", "predicate": "rel", "object": "B", "object-entity": True} {"type": "relationship", "subject": "A", "predicate": "rel", "object": "B", "object-entity": True}
] ]
triples, entity_contexts = agent_extractor.process_extraction_data(data, sample_metadata) triples, entity_contexts, _ = agent_extractor.process_extraction_data(data, sample_metadata)
# Should process valid items and ignore unknown types # Should process valid items and ignore unknown types
assert len(entity_contexts) == 1 # Only the definition creates entity context assert len(entity_contexts) == 1 # Only the definition creates entity context

View file

@ -168,7 +168,7 @@ class TestAgentKgExtractionEdgeCases:
"""Test processing with empty or minimal metadata""" """Test processing with empty or minimal metadata"""
# Test with None metadata - may not raise AttributeError depending on implementation # Test with None metadata - may not raise AttributeError depending on implementation
try: try:
triples, contexts = agent_extractor.process_extraction_data([], None) triples, contexts, _ = agent_extractor.process_extraction_data([], None)
# If it doesn't raise, check the results # If it doesn't raise, check the results
assert len(triples) == 0 assert len(triples) == 0
assert len(contexts) == 0 assert len(contexts) == 0
@ -178,14 +178,14 @@ class TestAgentKgExtractionEdgeCases:
# Test with metadata without ID # Test with metadata without ID
metadata = Metadata(id=None) metadata = Metadata(id=None)
triples, contexts = agent_extractor.process_extraction_data([], metadata) triples, contexts, _ = agent_extractor.process_extraction_data([], metadata)
assert len(triples) == 0 assert len(triples) == 0
assert len(contexts) == 0 assert len(contexts) == 0
# Test with metadata with empty string ID # Test with metadata with empty string ID
metadata = Metadata(id="") metadata = Metadata(id="")
data = [{"type": "definition", "entity": "Test", "definition": "Test def"}] data = [{"type": "definition", "entity": "Test", "definition": "Test def"}]
triples, contexts = agent_extractor.process_extraction_data(data, metadata) triples, contexts, _ = agent_extractor.process_extraction_data(data, metadata)
# Should not create subject-of triples when ID is empty string # Should not create subject-of triples when ID is empty string
subject_of_triples = [t for t in triples if t.p.iri == SUBJECT_OF] subject_of_triples = [t for t in triples if t.p.iri == SUBJECT_OF]
@ -213,7 +213,7 @@ class TestAgentKgExtractionEdgeCases:
for entity in special_entities for entity in special_entities
] ]
triples, contexts = agent_extractor.process_extraction_data(data, metadata) triples, contexts, _ = agent_extractor.process_extraction_data(data, metadata)
# Verify all entities were processed # Verify all entities were processed
assert len(contexts) == len(special_entities) assert len(contexts) == len(special_entities)
@ -234,7 +234,7 @@ class TestAgentKgExtractionEdgeCases:
{"type": "definition", "entity": "Test Entity", "definition": long_definition} {"type": "definition", "entity": "Test Entity", "definition": long_definition}
] ]
triples, contexts = agent_extractor.process_extraction_data(data, metadata) triples, contexts, _ = agent_extractor.process_extraction_data(data, metadata)
# Should handle long definitions without issues # Should handle long definitions without issues
assert len(contexts) == 1 assert len(contexts) == 1
@ -256,7 +256,7 @@ class TestAgentKgExtractionEdgeCases:
{"type": "definition", "entity": "AI", "definition": "Another AI definition"}, # Duplicate {"type": "definition", "entity": "AI", "definition": "Another AI definition"}, # Duplicate
] ]
triples, contexts = agent_extractor.process_extraction_data(data, metadata) triples, contexts, _ = agent_extractor.process_extraction_data(data, metadata)
# Should process all entries (including duplicates) # Should process all entries (including duplicates)
assert len(contexts) == 4 assert len(contexts) == 4
@ -280,7 +280,7 @@ class TestAgentKgExtractionEdgeCases:
{"type": "relationship", "subject": "test", "predicate": "test", "object": "", "object-entity": True}, {"type": "relationship", "subject": "test", "predicate": "test", "object": "", "object-entity": True},
] ]
triples, contexts = agent_extractor.process_extraction_data(data, metadata) triples, contexts, _ = agent_extractor.process_extraction_data(data, metadata)
# Should handle empty strings by creating URIs (even if empty) # Should handle empty strings by creating URIs (even if empty)
assert len(contexts) == 3 assert len(contexts) == 3
@ -306,7 +306,7 @@ class TestAgentKgExtractionEdgeCases:
} }
] ]
triples, contexts = agent_extractor.process_extraction_data(data, metadata) triples, contexts, _ = agent_extractor.process_extraction_data(data, metadata)
# Should handle JSON strings in definitions without parsing them # Should handle JSON strings in definitions without parsing them
assert len(contexts) == 2 assert len(contexts) == 2
@ -334,7 +334,7 @@ class TestAgentKgExtractionEdgeCases:
{"type": "relationship", "subject": "A", "predicate": "rel7", "object": "F", "object-entity": 1}, {"type": "relationship", "subject": "A", "predicate": "rel7", "object": "F", "object-entity": 1},
] ]
triples, contexts = agent_extractor.process_extraction_data(data, metadata) triples, contexts, _ = agent_extractor.process_extraction_data(data, metadata)
# Should process all relationships # Should process all relationships
# Note: The current implementation has some logic issues that these tests document # Note: The current implementation has some logic issues that these tests document
@ -416,7 +416,7 @@ class TestAgentKgExtractionEdgeCases:
import time import time
start_time = time.time() start_time = time.time()
triples, contexts = agent_extractor.process_extraction_data(large_data, metadata) triples, contexts, _ = agent_extractor.process_extraction_data(large_data, metadata)
end_time = time.time() end_time = time.time()
processing_time = end_time - start_time processing_time = end_time - start_time