diff --git a/tests/integration/test_load_structured_data_integration.py b/tests/integration/test_load_structured_data_integration.py
new file mode 100644
index 00000000..6d087f05
--- /dev/null
+++ b/tests/integration/test_load_structured_data_integration.py
@@ -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 = """
+
+
+
+ John Smith
+ john@email.com
+ 35
+ US
+ active
+
+
+ Jane Doe
+ jane@email.com
+ 28
+ CA
+ active
+
+
+ Bob Johnson
+ bob@company.org
+ 42
+ UK
+ inactive
+
+
+"""
+
+ 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)
\ No newline at end of file
diff --git a/tests/integration/test_load_structured_data_websocket.py b/tests/integration/test_load_structured_data_websocket.py
new file mode 100644
index 00000000..ab647f18
--- /dev/null
+++ b/tests/integration/test_load_structured_data_websocket.py
@@ -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)
\ No newline at end of file
diff --git a/tests/unit/test_cli/test_load_structured_data.py b/tests/unit/test_cli/test_load_structured_data.py
new file mode 100644
index 00000000..f84e1e5c
--- /dev/null
+++ b/tests/unit/test_cli/test_load_structured_data.py
@@ -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 = """
+
+
+
+ John Smith
+ john@email.com
+ 35
+
+
+ Jane Doe
+ jane@email.com
+ 28
+
+
+"""
+
+ 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 = """
+
+