apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-24 12:41:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Tests

Browse source
This commit is contained in:
Cyber MacGeddon 2025-09-05 16:10:35 +01:00
parent 0b7620bc04
commit 4c5a4b5e81
3 changed files with 1443 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

531
tests/integration/test_load_structured_data_integration.py Normal file
View file

@ -0,0 +1,531 @@
"""
Integration tests for tg-load-structured-data with actual TrustGraph instance.
Tests end-to-end functionality including WebSocket connections and data storage.
"""
import pytest
import asyncio
import json
import tempfile
import os
import csv
import time
from unittest.mock import Mock, patch, AsyncMock
from websockets.asyncio.client import connect
from trustgraph.cli.load_structured_data import load_structured_data
@pytest.mark.integration
class TestLoadStructuredDataIntegration:
"""Integration tests for complete pipeline"""
def setup_method(self):
"""Set up test fixtures"""
self.api_url = "http://localhost:8088"
self.test_schema_name = "integration_test_schema"
self.test_csv_data = """name,email,age,country,status
John Smith,john@email.com,35,US,active
Jane Doe,jane@email.com,28,CA,active
Bob Johnson,bob@company.org,42,UK,inactive
Alice Brown,alice@email.com,31,AU,active
Charlie Davis,charlie@email.com,39,DE,inactive"""
self.test_json_data = [
{"name": "John Smith", "email": "john@email.com", "age": 35, "country": "US", "status": "active"},
{"name": "Jane Doe", "email": "jane@email.com", "age": 28, "country": "CA", "status": "active"},
{"name": "Bob Johnson", "email": "bob@company.org", "age": 42, "country": "UK", "status": "inactive"}
]
self.test_xml_data = """<?xml version="1.0"?>
<ROOT>
<data>
<record>
<field name="name">John Smith</field>
<field name="email">john@email.com</field>
<field name="age">35</field>
<field name="country">US</field>
<field name="status">active</field>
</record>
<record>
<field name="name">Jane Doe</field>
<field name="email">jane@email.com</field>
<field name="age">28</field>
<field name="country">CA</field>
<field name="status">active</field>
</record>
<record>
<field name="name">Bob Johnson</field>
<field name="email">bob@company.org</field>
<field name="age">42</field>
<field name="country">UK</field>
<field name="status">inactive</field>
</record>
</data>
</ROOT>"""
self.test_descriptor = {
"version": "1.0",
"metadata": {
"name": "IntegrationTest",
"description": "Test descriptor for integration tests",
"author": "Test Suite"
},
"format": {
"type": "csv",
"encoding": "utf-8",
"options": {
"header": True,
"delimiter": ","
}
},
"mappings": [
{
"source_field": "name",
"target_field": "name",
"transforms": [{"type": "trim"}],
"validation": [{"type": "required"}]
},
{
"source_field": "email",
"target_field": "email",
"transforms": [{"type": "trim"}, {"type": "lower"}],
"validation": [{"type": "required"}]
},
{
"source_field": "age",
"target_field": "age",
"transforms": [{"type": "to_int"}],
"validation": [{"type": "required"}]
},
{
"source_field": "country",
"target_field": "country",
"transforms": [{"type": "trim"}, {"type": "upper"}],
"validation": [{"type": "required"}]
},
{
"source_field": "status",
"target_field": "status",
"transforms": [{"type": "trim"}, {"type": "lower"}],
"validation": [{"type": "required"}]
}
],
"output": {
"format": "trustgraph-objects",
"schema_name": self.test_schema_name,
"options": {
"confidence": 0.9,
"batch_size": 3
}
}
}
def create_temp_file(self, content, suffix='.txt'):
"""Create a temporary file with given content"""
temp_file = tempfile.NamedTemporaryFile(mode='w', suffix=suffix, delete=False)
temp_file.write(content)
temp_file.flush()
temp_file.close()
return temp_file.name
def cleanup_temp_file(self, file_path):
"""Clean up temporary file"""
try:
os.unlink(file_path)
except:
pass
# End-to-end Pipeline Tests
@pytest.mark.asyncio
async def test_csv_to_trustgraph_pipeline(self):
"""Test complete CSV to TrustGraph pipeline"""
input_file = self.create_temp_file(self.test_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
# Test with dry run first
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
dry_run=True,
batch_size=2,
flow='obj-ex'
)
# Should complete without errors in dry run mode
assert result is None # dry_run returns None
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
@pytest.mark.asyncio
async def test_xml_to_trustgraph_pipeline(self):
"""Test complete XML to TrustGraph pipeline"""
# Create XML descriptor
xml_descriptor = {
**self.test_descriptor,
"format": {
"type": "xml",
"encoding": "utf-8",
"options": {
"record_path": "/ROOT/data/record",
"field_attribute": "name"
}
}
}
input_file = self.create_temp_file(self.test_xml_data, '.xml')
descriptor_file = self.create_temp_file(json.dumps(xml_descriptor), '.json')
try:
# Test with dry run
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
dry_run=True,
batch_size=2,
flow='obj-ex'
)
assert result is None # dry_run returns None
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
@pytest.mark.asyncio
async def test_json_to_trustgraph_pipeline(self):
"""Test complete JSON to TrustGraph pipeline"""
json_descriptor = {
**self.test_descriptor,
"format": {
"type": "json",
"encoding": "utf-8"
}
}
input_file = self.create_temp_file(json.dumps(self.test_json_data), '.json')
descriptor_file = self.create_temp_file(json.dumps(json_descriptor), '.json')
try:
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
dry_run=True,
batch_size=2,
flow='obj-ex'
)
assert result is None # dry_run returns None
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
# Batching Integration Tests
@pytest.mark.asyncio
async def test_large_dataset_batching(self):
"""Test batching with larger dataset"""
# Generate larger dataset
large_csv_data = "name,email,age,country,status\n"
for i in range(1000):
large_csv_data += f"User{i},user{i}@example.com,{25+i%40},US,active\n"
input_file = self.create_temp_file(large_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
start_time = time.time()
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
dry_run=True,
batch_size=50, # Test with moderate batch size
flow='obj-ex'
)
end_time = time.time()
processing_time = end_time - start_time
# Should process 1000 records reasonably quickly
assert processing_time < 30 # Should complete in under 30 seconds
assert result is None # dry_run returns None
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
@pytest.mark.asyncio
async def test_batch_size_performance(self):
"""Test different batch sizes for performance"""
# Generate test dataset
test_csv_data = "name,email,age,country,status\n"
for i in range(100):
test_csv_data += f"User{i},user{i}@example.com,{25+i%40},US,active\n"
input_file = self.create_temp_file(test_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
# Test different batch sizes
batch_sizes = [1, 10, 25, 50, 100]
processing_times = {}
for batch_size in batch_sizes:
start_time = time.time()
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
dry_run=True,
batch_size=batch_size,
flow='obj-ex'
)
end_time = time.time()
processing_times[batch_size] = end_time - start_time
assert result is None # dry_run returns None
# All batch sizes should complete reasonably quickly
for batch_size, time_taken in processing_times.items():
assert time_taken < 10, f"Batch size {batch_size} took {time_taken}s"
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
# Parse-Only Mode Tests
@pytest.mark.asyncio
async def test_parse_only_mode(self):
"""Test parse-only mode functionality"""
input_file = self.create_temp_file(self.test_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
output_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
output_file.close()
try:
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
parse_only=True,
output_file=output_file.name
)
# Check output file was created and contains parsed data
assert os.path.exists(output_file.name)
with open(output_file.name, 'r') as f:
parsed_data = json.load(f)
assert isinstance(parsed_data, list)
assert len(parsed_data) == 5 # Should have 5 records
assert parsed_data[0]["name"] == "John Smith"
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
self.cleanup_temp_file(output_file.name)
# Schema Suggestion Integration Tests
@patch('trustgraph.cli.load_structured_data.TrustGraphAPI')
@pytest.mark.asyncio
async def test_schema_suggestion_integration(self, mock_api_class):
"""Test schema suggestion integration with API"""
# Setup mock API responses
mock_api = Mock()
mock_api_class.return_value = mock_api
mock_config_api = Mock()
mock_api.config.return_value = mock_config_api
mock_config_api.get_config_items.return_value = {
"schema": {
"customer": '{"name": "customer", "description": "Customer records"}',
"product": '{"name": "product", "description": "Product catalog"}'
}
}
mock_flow = Mock()
mock_api.flow.return_value = mock_flow
mock_flow.id.return_value = mock_flow
mock_prompt_client = Mock()
mock_flow.prompt.return_value = mock_prompt_client
mock_prompt_client.schema_selection.return_value = "The customer schema is most appropriate for this data containing names and emails."
input_file = self.create_temp_file(self.test_csv_data, '.csv')
try:
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
suggest_schema=True,
sample_size=100,
sample_chars=500
)
# Verify API calls were made correctly
mock_config_api.get_config_items.assert_called_once()
mock_prompt_client.schema_selection.assert_called_once()
# Check that schemas were passed correctly
call_args = mock_prompt_client.schema_selection.call_args
assert 'schemas' in call_args.kwargs
assert 'sample' in call_args.kwargs
finally:
self.cleanup_temp_file(input_file)
# Descriptor Generation Integration Tests
@patch('trustgraph.cli.load_structured_data.TrustGraphAPI')
@pytest.mark.asyncio
async def test_descriptor_generation_integration(self, mock_api_class):
"""Test descriptor generation integration"""
# Setup mock API
mock_api = Mock()
mock_api_class.return_value = mock_api
mock_config_api = Mock()
mock_api.config.return_value = mock_config_api
mock_config_api.get_config_items.return_value = {
"schema": {
"customer": '{"name": "customer", "fields": [{"name": "name", "type": "string"}]}'
}
}
mock_flow = Mock()
mock_api.flow.return_value = mock_flow
mock_flow.id.return_value = mock_flow
mock_prompt_client = Mock()
mock_flow.prompt.return_value = mock_prompt_client
# Mock descriptor generation response
generated_descriptor = {**self.test_descriptor}
mock_prompt_client.diagnose_structured_data.return_value = json.dumps(generated_descriptor)
input_file = self.create_temp_file(self.test_csv_data, '.csv')
try:
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
generate_descriptor=True,
sample_chars=1000
)
# Verify API calls
mock_prompt_client.diagnose_structured_data.assert_called_once()
# Check call arguments
call_args = mock_prompt_client.diagnose_structured_data.call_args
assert 'schemas' in call_args.kwargs
assert 'sample' in call_args.kwargs
finally:
self.cleanup_temp_file(input_file)
# Error Handling Integration Tests
@pytest.mark.asyncio
async def test_malformed_data_handling(self):
"""Test handling of malformed data"""
malformed_csv = """name,email,age
John Smith,john@email.com,35
Jane Doe,jane@email.com # Missing age field
Bob Johnson,bob@company.org,not_a_number"""
input_file = self.create_temp_file(malformed_csv, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
# Should handle malformed data gracefully
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
dry_run=True,
max_errors=5 # Allow some errors
)
# Should complete even with some malformed records
assert result is None # dry_run returns None
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
# WebSocket Connection Tests
@pytest.mark.asyncio
async def test_websocket_connection_handling(self):
"""Test WebSocket connection behavior"""
input_file = self.create_temp_file(self.test_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
# Test with invalid API URL (should fail gracefully)
with pytest.raises(Exception): # Connection error expected
result = load_structured_data(
api_url="http://invalid-url:9999",
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=2,
flow='obj-ex'
)
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
# Flow Parameter Tests
@pytest.mark.asyncio
async def test_flow_parameter_integration(self):
"""Test flow parameter functionality"""
input_file = self.create_temp_file(self.test_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
# Test with different flow values
flows = ['default', 'obj-ex', 'custom-flow']
for flow in flows:
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
dry_run=True,
flow=flow
)
assert result is None # dry_run returns None
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
# Mixed Format Tests
@pytest.mark.asyncio
async def test_encoding_variations(self):
"""Test different encoding variations"""
# Test UTF-8 with BOM
utf8_bom_data = '\ufeff' + self.test_csv_data
input_file = self.create_temp_file(utf8_bom_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
dry_run=True
)
assert result is None # Should handle BOM correctly
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)

481
tests/integration/test_load_structured_data_websocket.py Normal file
View file

@ -0,0 +1,481 @@
"""
WebSocket-specific integration tests for tg-load-structured-data.
Tests WebSocket connection handling, message formats, and batching behavior.
"""
import pytest
import asyncio
import json
import tempfile
import os
from unittest.mock import Mock, patch, AsyncMock, MagicMock
import websockets
from websockets.exceptions import ConnectionClosedError, InvalidHandshake
from trustgraph.cli.load_structured_data import load_structured_data
@pytest.mark.integration
class TestLoadStructuredDataWebSocket:
"""WebSocket-specific integration tests"""
def setup_method(self):
"""Set up test fixtures"""
self.api_url = "http://localhost:8088"
self.ws_url = "ws://localhost:8088"
self.test_csv_data = """name,email,age,country
John Smith,john@email.com,35,US
Jane Doe,jane@email.com,28,CA
Bob Johnson,bob@company.org,42,UK
Alice Brown,alice@email.com,31,AU
Charlie Davis,charlie@email.com,39,DE"""
self.test_descriptor = {
"version": "1.0",
"format": {
"type": "csv",
"encoding": "utf-8",
"options": {"header": True, "delimiter": ","}
},
"mappings": [
{"source_field": "name", "target_field": "name", "transforms": [{"type": "trim"}]},
{"source_field": "email", "target_field": "email", "transforms": [{"type": "lower"}]},
{"source_field": "age", "target_field": "age", "transforms": [{"type": "to_int"}]},
{"source_field": "country", "target_field": "country", "transforms": [{"type": "upper"}]}
],
"output": {
"format": "trustgraph-objects",
"schema_name": "test_customer",
"options": {"confidence": 0.9, "batch_size": 2}
}
}
def create_temp_file(self, content, suffix='.txt'):
"""Create a temporary file with given content"""
temp_file = tempfile.NamedTemporaryFile(mode='w', suffix=suffix, delete=False)
temp_file.write(content)
temp_file.flush()
temp_file.close()
return temp_file.name
def cleanup_temp_file(self, file_path):
"""Clean up temporary file"""
try:
os.unlink(file_path)
except:
pass
@pytest.mark.asyncio
async def test_websocket_message_format(self):
"""Test that WebSocket messages are formatted correctly for batching"""
messages_sent = []
# Mock WebSocket connection
async def mock_websocket_handler(websocket, path):
try:
while True:
message = await websocket.recv()
messages_sent.append(json.loads(message))
except websockets.exceptions.ConnectionClosed:
pass
# Start mock WebSocket server
server = await websockets.serve(mock_websocket_handler, "localhost", 8089)
try:
input_file = self.create_temp_file(self.test_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
# Test with mock server
with patch('websockets.asyncio.client.connect') as mock_connect:
mock_ws = AsyncMock()
mock_connect.return_value.__aenter__.return_value = mock_ws
# Capture messages sent
sent_messages = []
mock_ws.send = AsyncMock(side_effect=lambda msg: sent_messages.append(json.loads(msg)))
try:
result = load_structured_data(
api_url="http://localhost:8089",
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=2,
flow='obj-ex'
)
# Verify message format
assert len(sent_messages) > 0
for message in sent_messages:
# Check required fields
assert "metadata" in message
assert "schema_name" in message
assert "values" in message
assert "confidence" in message
assert "source_span" in message
# Check metadata structure
metadata = message["metadata"]
assert "id" in metadata
assert "metadata" in metadata
assert "user" in metadata
assert "collection" in metadata
# Check batched values format
values = message["values"]
assert isinstance(values, list), "Values should be a list (batched)"
assert len(values) <= 2, "Batch size should be respected"
# Check each object in batch
for obj in values:
assert isinstance(obj, dict)
assert "name" in obj
assert "email" in obj
assert "age" in obj
assert "country" in obj
# Check transformations were applied
assert obj["email"].islower(), "Email should be lowercase"
assert obj["country"].isupper(), "Country should be uppercase"
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
finally:
server.close()
await server.wait_closed()
@pytest.mark.asyncio
async def test_websocket_connection_retry(self):
"""Test WebSocket connection retry behavior"""
input_file = self.create_temp_file(self.test_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
# Test connection to non-existent server
with pytest.raises(Exception) as exc_info:
result = load_structured_data(
api_url="http://localhost:9999", # Non-existent server
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=2,
flow='obj-ex'
)
# Should get connection error
assert "connection" in str(exc_info.value).lower() or "refused" in str(exc_info.value).lower()
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
@pytest.mark.asyncio
async def test_websocket_large_message_handling(self):
"""Test WebSocket handling of large batched messages"""
# Generate larger dataset
large_csv_data = "name,email,age,country\n"
for i in range(100):
large_csv_data += f"User{i},user{i}@example.com,{25+i%40},US\n"
# Create descriptor with larger batch size
large_batch_descriptor = {
**self.test_descriptor,
"output": {
**self.test_descriptor["output"],
"batch_size": 50 # Large batch size
}
}
input_file = self.create_temp_file(large_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(large_batch_descriptor), '.json')
try:
with patch('websockets.asyncio.client.connect') as mock_connect:
mock_ws = AsyncMock()
mock_connect.return_value.__aenter__.return_value = mock_ws
sent_messages = []
mock_ws.send = AsyncMock(side_effect=lambda msg: sent_messages.append(json.loads(msg)))
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=50,
flow='obj-ex'
)
# Should handle large batches
assert len(sent_messages) >= 2 # 100 records with batch_size=50 -> 2 messages
# Check message sizes
for message in sent_messages:
values = message["values"]
assert len(values) <= 50
# Check message is not too large (rough size check)
message_size = len(json.dumps(message))
assert message_size < 1024 * 1024 # Less than 1MB per message
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
@pytest.mark.asyncio
async def test_websocket_connection_interruption(self):
"""Test handling of WebSocket connection interruptions"""
input_file = self.create_temp_file(self.test_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
with patch('websockets.asyncio.client.connect') as mock_connect:
mock_ws = AsyncMock()
mock_connect.return_value.__aenter__.return_value = mock_ws
# Simulate connection being closed mid-send
call_count = 0
def send_with_failure(msg):
nonlocal call_count
call_count += 1
if call_count > 1: # Fail after first message
raise ConnectionClosedError(None, None)
return AsyncMock()
mock_ws.send.side_effect = send_with_failure
# Should handle connection errors
with pytest.raises(ConnectionClosedError):
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=2,
flow='obj-ex'
)
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
@pytest.mark.asyncio
async def test_websocket_url_conversion(self):
"""Test proper URL conversion from HTTP to WebSocket"""
input_file = self.create_temp_file(self.test_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
with patch('websockets.asyncio.client.connect') as mock_connect:
mock_ws = AsyncMock()
mock_connect.return_value.__aenter__.return_value = mock_ws
mock_ws.send = AsyncMock()
# Test HTTP URL conversion
result = load_structured_data(
api_url="http://localhost:8088", # HTTP URL
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=2,
flow='obj-ex'
)
# Check that WebSocket URL was used
mock_connect.assert_called_once()
called_url = mock_connect.call_args[0][0]
assert called_url.startswith("ws://")
assert "api/v1/flow/obj-ex/import/objects" in called_url
# Test HTTPS URL conversion
mock_connect.reset_mock()
result = load_structured_data(
api_url="https://example.com:8088", # HTTPS URL
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=2,
flow='test-flow'
)
# Check that secure WebSocket URL was used
mock_connect.assert_called_once()
called_url = mock_connect.call_args[0][0]
assert called_url.startswith("wss://")
assert "api/v1/flow/test-flow/import/objects" in called_url
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
@pytest.mark.asyncio
async def test_websocket_batch_ordering(self):
"""Test that batches are sent in correct order"""
# Create ordered test data
ordered_csv_data = "name,id\n"
for i in range(10):
ordered_csv_data += f"User{i:02d},{i}\n"
input_file = self.create_temp_file(ordered_csv_data, '.csv')
# Create descriptor for this test
ordered_descriptor = {
**self.test_descriptor,
"mappings": [
{"source_field": "name", "target_field": "name", "transforms": []},
{"source_field": "id", "target_field": "id", "transforms": [{"type": "to_int"}]}
],
"output": {
**self.test_descriptor["output"],
"batch_size": 3
}
}
descriptor_file = self.create_temp_file(json.dumps(ordered_descriptor), '.json')
try:
with patch('websockets.asyncio.client.connect') as mock_connect:
mock_ws = AsyncMock()
mock_connect.return_value.__aenter__.return_value = mock_ws
sent_messages = []
mock_ws.send = AsyncMock(side_effect=lambda msg: sent_messages.append(json.loads(msg)))
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=3,
flow='obj-ex'
)
# Should have 4 messages (10 records, batch_size=3: 3+3+3+1)
assert len(sent_messages) == 4
# Check ordering within batches
expected_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
actual_ids = []
for message in sent_messages:
values = message["values"]
for obj in values:
actual_ids.append(int(obj["id"]))
assert actual_ids == expected_ids, "Records should maintain order"
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
@pytest.mark.asyncio
async def test_websocket_authentication_headers(self):
"""Test WebSocket connection with authentication headers"""
input_file = self.create_temp_file(self.test_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
with patch('websockets.asyncio.client.connect') as mock_connect:
mock_ws = AsyncMock()
mock_connect.return_value.__aenter__.return_value = mock_ws
mock_ws.send = AsyncMock()
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=2,
flow='obj-ex'
)
# Verify WebSocket connect was called
mock_connect.assert_called_once()
# In real implementation, could check for auth headers
# For now, just verify the connection was attempted
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
@pytest.mark.asyncio
async def test_websocket_empty_batch_handling(self):
"""Test handling of empty batches"""
# Create CSV with some invalid records
invalid_csv_data = """name,email,age,country
,invalid@email,not_a_number,
Valid User,valid@email.com,25,US"""
input_file = self.create_temp_file(invalid_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
with patch('websockets.asyncio.client.connect') as mock_connect:
mock_ws = AsyncMock()
mock_connect.return_value.__aenter__.return_value = mock_ws
sent_messages = []
mock_ws.send = AsyncMock(side_effect=lambda msg: sent_messages.append(json.loads(msg)))
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=2,
flow='obj-ex'
)
# Should still send messages for valid records
assert len(sent_messages) >= 1
# Check that messages are not empty
for message in sent_messages:
values = message["values"]
assert len(values) > 0, "Should not send empty batches"
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
@pytest.mark.asyncio
async def test_websocket_progress_reporting(self):
"""Test progress reporting during WebSocket sends"""
# Generate larger dataset for progress testing
progress_csv_data = "name,email,age\n"
for i in range(50):
progress_csv_data += f"User{i},user{i}@example.com,{25+i}\n"
input_file = self.create_temp_file(progress_csv_data, '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.test_descriptor), '.json')
try:
with patch('websockets.asyncio.client.connect') as mock_connect:
mock_ws = AsyncMock()
mock_connect.return_value.__aenter__.return_value = mock_ws
send_count = 0
def count_sends(msg):
nonlocal send_count
send_count += 1
return AsyncMock()
mock_ws.send.side_effect = count_sends
# Capture logging output to check for progress messages
with patch('logging.getLogger') as mock_logger:
mock_log = Mock()
mock_logger.return_value = mock_log
result = load_structured_data(
api_url=self.api_url,
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=10, # Should result in 5 batches
flow='obj-ex',
verbose=True
)
# Should have sent multiple batches
assert send_count >= 5
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)

431
tests/unit/test_cli/test_load_structured_data.py Normal file
View file

@ -0,0 +1,431 @@
"""
Unit tests for tg-load-structured-data CLI command.
Tests all modes: suggest-schema, generate-descriptor, parse-only, full pipeline.
"""
import pytest
import json
import tempfile
import os
import csv
import xml.etree.ElementTree as ET
from unittest.mock import Mock, patch, AsyncMock, MagicMock, call
from io import StringIO
import asyncio
# Import the function we're testing
from trustgraph.cli.load_structured_data import (
load_structured_data,
parse_csv_data,
parse_json_data,
parse_xml_data,
apply_transformations
)
class TestLoadStructuredDataUnit:
"""Unit tests for load_structured_data functionality"""
def setup_method(self):
"""Set up test fixtures"""
self.test_csv_data = """name,email,age,country
John Smith,john@email.com,35,US
Jane Doe,jane@email.com,28,CA
Bob Johnson,bob@company.org,42,UK"""
self.test_json_data = [
{"name": "John Smith", "email": "john@email.com", "age": 35, "country": "US"},
{"name": "Jane Doe", "email": "jane@email.com", "age": 28, "country": "CA"}
]
self.test_xml_data = """<?xml version="1.0"?>
<ROOT>
<data>
<record>
<field name="name">John Smith</field>
<field name="email">john@email.com</field>
<field name="age">35</field>
</record>
<record>
<field name="name">Jane Doe</field>
<field name="email">jane@email.com</field>
<field name="age">28</field>
</record>
</data>
</ROOT>"""
self.test_descriptor = {
"version": "1.0",
"format": {"type": "csv", "encoding": "utf-8", "options": {"header": True}},
"mappings": [
{"source_field": "name", "target_field": "name", "transforms": [{"type": "trim"}]},
{"source_field": "email", "target_field": "email", "transforms": [{"type": "lower"}]}
],
"output": {
"format": "trustgraph-objects",
"schema_name": "customer",
"options": {"confidence": 0.9, "batch_size": 100}
}
}
# CSV Parsing Tests
def test_csv_parsing_with_header(self):
"""Test CSV parsing with header row"""
format_info = {"type": "csv", "encoding": "utf-8", "options": {"header": True, "delimiter": ","}}
records = parse_csv_data(self.test_csv_data, format_info)
assert len(records) == 3
assert records[0]["name"] == "John Smith"
assert records[0]["email"] == "john@email.com"
assert records[1]["country"] == "CA"
def test_csv_parsing_without_header(self):
"""Test CSV parsing without header row"""
csv_data = """John Smith,john@email.com,35,US
Jane Doe,jane@email.com,28,CA"""
format_info = {"type": "csv", "encoding": "utf-8", "options": {"header": False, "delimiter": ","}}
records = parse_csv_data(csv_data, format_info)
assert len(records) == 2
# Should use column indices as keys
assert records[0]["0"] == "John Smith"
assert records[0]["1"] == "john@email.com"
def test_csv_parsing_custom_delimiter(self):
"""Test CSV parsing with custom delimiter"""
csv_data = """name;email;age
John Smith;john@email.com;35
Jane Doe;jane@email.com;28"""
format_info = {"type": "csv", "encoding": "utf-8", "options": {"header": True, "delimiter": ";"}}
records = parse_csv_data(csv_data, format_info)
assert len(records) == 2
assert records[0]["name"] == "John Smith"
assert records[1]["email"] == "jane@email.com"
# JSON Parsing Tests
def test_json_parsing_array_format(self):
"""Test JSON parsing with array format"""
json_data = json.dumps(self.test_json_data)
format_info = {"type": "json", "encoding": "utf-8"}
records = parse_json_data(json_data, format_info)
assert len(records) == 2
assert records[0]["name"] == "John Smith"
assert records[1]["country"] == "CA"
def test_json_parsing_newline_delimited(self):
"""Test JSON parsing with newline-delimited format"""
json_data = '{"name": "John", "age": 35}\n{"name": "Jane", "age": 28}'
format_info = {"type": "json", "encoding": "utf-8", "options": {"newline_delimited": True}}
records = parse_json_data(json_data, format_info)
assert len(records) == 2
assert records[0]["name"] == "John"
assert records[1]["age"] == 28
def test_json_parsing_malformed_data(self):
"""Test JSON parsing with malformed data"""
json_data = '{"name": "John", "age": 35' # Missing closing brace
format_info = {"type": "json", "encoding": "utf-8"}
with pytest.raises(json.JSONDecodeError):
parse_json_data(json_data, format_info)
# XML Parsing Tests
def test_xml_parsing_xpath_expressions(self):
"""Test XML parsing with XPath expressions"""
format_info = {
"type": "xml",
"encoding": "utf-8",
"options": {
"record_path": "/ROOT/data/record",
"field_attribute": "name"
}
}
records = parse_xml_data(self.test_xml_data, format_info)
assert len(records) == 2
assert records[0]["name"] == "John Smith"
assert records[0]["email"] == "john@email.com"
assert records[1]["name"] == "Jane Doe"
def test_xml_parsing_field_attributes(self):
"""Test XML parsing with field attributes"""
xml_data = """<?xml version="1.0"?>
<records>
<item id="1" name="Product A" price="19.99"/>
<item id="2" name="Product B" price="29.99"/>
</records>"""
format_info = {
"type": "xml",
"encoding": "utf-8",
"options": {"record_path": "//item"}
}
records = parse_xml_data(xml_data, format_info)
assert len(records) == 2
assert records[0]["id"] == "1"
assert records[0]["name"] == "Product A"
assert records[1]["price"] == "29.99"
# Data Transformation Tests
def test_field_mappings_application(self):
"""Test field mapping application"""
record = {"source_name": "John Smith", "source_email": "JOHN@EMAIL.COM"}
mappings = [
{"source_field": "source_name", "target_field": "name", "transforms": []},
{"source_field": "source_email", "target_field": "email", "transforms": []}
]
result = apply_transformations(record, mappings)
assert result["name"] == "John Smith"
assert result["email"] == "JOHN@EMAIL.COM"
def test_data_transforms_trim(self):
"""Test trim transformation"""
record = {"name": " John Smith ", "email": "john@email.com"}
mappings = [
{"source_field": "name", "target_field": "name", "transforms": [{"type": "trim"}]}
]
result = apply_transformations(record, mappings)
assert result["name"] == "John Smith"
def test_data_transforms_case_conversion(self):
"""Test case conversion transformations"""
record = {"name": "John Smith", "email": "JOHN@EMAIL.COM"}
mappings = [
{"source_field": "name", "target_field": "name", "transforms": [{"type": "upper"}]},
{"source_field": "email", "target_field": "email", "transforms": [{"type": "lower"}]}
]
result = apply_transformations(record, mappings)
assert result["name"] == "JOHN SMITH"
assert result["email"] == "john@email.com"
def test_data_transforms_type_conversion(self):
"""Test type conversion transformations"""
record = {"age": "35", "price": "19.99", "active": "true"}
mappings = [
{"source_field": "age", "target_field": "age", "transforms": [{"type": "to_int"}]},
{"source_field": "price", "target_field": "price", "transforms": [{"type": "to_float"}]},
{"source_field": "active", "target_field": "active", "transforms": [{"type": "to_bool"}]}
]
result = apply_transformations(record, mappings)
# All values should be converted to strings for ExtractedObject compatibility
assert result["age"] == "35"
assert result["price"] == "19.99"
assert result["active"] == "True"
# Batching Functionality Tests
def test_batch_processing_single_object(self):
"""Test batch processing with single object"""
records = [{"name": "John", "age": "35"}]
batch_size = 10
schema_name = "customer"
# Simulate batching logic
batches = []
for i in range(0, len(records), batch_size):
batch_records = records[i:i + batch_size]
batch_values = [record for record in batch_records]
batches.append({
"schema_name": schema_name,
"values": batch_values
})
assert len(batches) == 1
assert len(batches[0]["values"]) == 1
assert batches[0]["values"][0]["name"] == "John"
def test_batch_processing_multiple_objects(self):
"""Test batch processing with multiple objects"""
records = [{"name": f"User{i}", "age": str(20+i)} for i in range(5)]
batch_size = 2
schema_name = "customer"
# Simulate batching logic
batches = []
for i in range(0, len(records), batch_size):
batch_records = records[i:i + batch_size]
batch_values = [record for record in batch_records]
batches.append({
"schema_name": schema_name,
"values": batch_values
})
assert len(batches) == 3 # 5 records with batch_size=2 -> 3 batches
assert len(batches[0]["values"]) == 2
assert len(batches[1]["values"]) == 2
assert len(batches[2]["values"]) == 1
assert batches[2]["values"][0]["name"] == "User4"
def test_batch_size_configuration(self):
"""Test batch size configuration"""
records = [{"id": str(i)} for i in range(10)]
for batch_size in [1, 3, 5, 10, 15]:
batches = []
for i in range(0, len(records), batch_size):
batch_records = records[i:i + batch_size]
batches.append(batch_records)
expected_batches = (len(records) + batch_size - 1) // batch_size
assert len(batches) == expected_batches
# Check all records are included
total_records = sum(len(batch) for batch in batches)
assert total_records == len(records)
# Schema Suggestion Tests
@patch('trustgraph.cli.load_structured_data.TrustGraphAPI')
def test_suggest_schema_api_integration(self, mock_api_class):
"""Test schema suggestion with API integration"""
# Setup mock API
mock_api = Mock()
mock_api_class.return_value = mock_api
mock_config_api = Mock()
mock_api.config.return_value = mock_config_api
mock_config_api.get_config_items.return_value = {"schema": {"customer": '{"name": "customer"}'}}
mock_flow = Mock()
mock_api.flow.return_value = mock_flow
mock_flow.id.return_value = mock_flow
mock_prompt_client = Mock()
mock_flow.prompt.return_value = mock_prompt_client
mock_prompt_client.schema_selection.return_value = "customer schema looks best for this data"
with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as f:
f.write(self.test_csv_data)
f.flush()
try:
# This should not raise an exception
result = load_structured_data(
api_url="http://localhost:8088",
input_file=f.name,
suggest_schema=True,
sample_size=100,
sample_chars=500
)
# Verify API calls were made
mock_config_api.get_config_items.assert_called_once()
mock_prompt_client.schema_selection.assert_called_once()
finally:
os.unlink(f.name)
# Descriptor Generation Tests
@patch('trustgraph.cli.load_structured_data.TrustGraphAPI')
def test_generate_descriptor_csv_format(self, mock_api_class):
"""Test descriptor generation for CSV format"""
# Setup mock API
mock_api = Mock()
mock_api_class.return_value = mock_api
mock_config_api = Mock()
mock_api.config.return_value = mock_config_api
mock_config_api.get_config_items.return_value = {"schema": {"customer": '{"name": "customer"}'}}
mock_flow = Mock()
mock_api.flow.return_value = mock_flow
mock_flow.id.return_value = mock_flow
mock_prompt_client = Mock()
mock_flow.prompt.return_value = mock_prompt_client
mock_prompt_client.diagnose_structured_data.return_value = json.dumps(self.test_descriptor)
with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as f:
f.write(self.test_csv_data)
f.flush()
try:
result = load_structured_data(
api_url="http://localhost:8088",
input_file=f.name,
generate_descriptor=True,
sample_chars=500
)
# Verify API calls were made
mock_prompt_client.diagnose_structured_data.assert_called_once()
finally:
os.unlink(f.name)
# Error Handling Tests
def test_file_not_found_error(self):
"""Test handling of file not found error"""
with pytest.raises(FileNotFoundError):
load_structured_data(
api_url="http://localhost:8088",
input_file="/nonexistent/file.csv"
)
def test_invalid_descriptor_format(self):
"""Test handling of invalid descriptor format"""
with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as input_file:
input_file.write(self.test_csv_data)
input_file.flush()
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as desc_file:
desc_file.write('{"invalid": "descriptor"}') # Missing required fields
desc_file.flush()
try:
with pytest.raises((KeyError, ValueError)):
load_structured_data(
api_url="http://localhost:8088",
input_file=input_file.name,
descriptor_file=desc_file.name
)
finally:
os.unlink(input_file.name)
os.unlink(desc_file.name)
def test_parsing_errors_handling(self):
"""Test handling of parsing errors"""
invalid_csv = "name,email\n\"unclosed quote,test@email.com"
format_info = {"type": "csv", "encoding": "utf-8", "options": {"header": True}}
# Should handle parsing errors gracefully
with pytest.raises(Exception):
parse_csv_data(invalid_csv, format_info)
# Validation Tests
def test_validation_rules_required_fields(self):
"""Test validation rules for required fields"""
record = {"name": "John", "email": ""} # Missing required email
mappings = [
{
"source_field": "name",
"target_field": "name",
"transforms": [],
"validation": [{"type": "required"}]
},
{
"source_field": "email",
"target_field": "email",
"transforms": [],
"validation": [{"type": "required"}]
}
]
result = apply_transformations(record, mappings)
# Should still process but may log warnings
assert "name" in result
assert "email" in result