mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-06-22 05:08:06 +02:00
IAM tech spec: Auth and access management current state and proposed
changes. Support for separate workspaces Addition of workspace CLI support for test purposes
This commit is contained in:
Cyber MacGeddon
parent
48da6c5f8b
commit
db05427d0e
219 changed files with 4875 additions and 2616 deletions
|
|
@ -58,7 +58,7 @@ class TestAgentStructuredQueryIntegration:
|
|||
async def test_agent_structured_query_basic_integration(self, agent_processor, structured_query_tool_config):
|
||||
"""Test basic agent integration with structured query tool"""
|
||||
# Arrange - Load tool configuration
|
||||
await agent_processor.on_tools_config(structured_query_tool_config, "v1")
|
||||
await agent_processor.on_tools_config("default", structured_query_tool_config, "v1")
|
||||
|
||||
# Create agent request
|
||||
request = AgentRequest(
|
||||
|
|
@ -119,6 +119,7 @@ Args: {
|
|||
# Mock flow parameter in agent_processor.on_request
|
||||
flow = MagicMock()
|
||||
flow.side_effect = flow_context
|
||||
flow.workspace = "default"
|
||||
|
||||
# Act
|
||||
await agent_processor.on_request(msg, consumer, flow)
|
||||
|
|
@ -146,7 +147,7 @@ Args: {
|
|||
async def test_agent_structured_query_error_handling(self, agent_processor, structured_query_tool_config):
|
||||
"""Test agent handling of structured query errors"""
|
||||
# Arrange
|
||||
await agent_processor.on_tools_config(structured_query_tool_config, "v1")
|
||||
await agent_processor.on_tools_config("default", structured_query_tool_config, "v1")
|
||||
|
||||
request = AgentRequest(
|
||||
question="Find data from a table that doesn't exist using structured query.",
|
||||
|
|
@ -199,6 +200,7 @@ Args: {
|
|||
|
||||
flow = MagicMock()
|
||||
flow.side_effect = flow_context
|
||||
flow.workspace = "default"
|
||||
|
||||
# Act
|
||||
await agent_processor.on_request(msg, consumer, flow)
|
||||
|
|
@ -221,7 +223,7 @@ Args: {
|
|||
async def test_agent_multi_step_structured_query_reasoning(self, agent_processor, structured_query_tool_config):
|
||||
"""Test agent using structured query in multi-step reasoning"""
|
||||
# Arrange
|
||||
await agent_processor.on_tools_config(structured_query_tool_config, "v1")
|
||||
await agent_processor.on_tools_config("default", structured_query_tool_config, "v1")
|
||||
|
||||
request = AgentRequest(
|
||||
question="First find all customers from California, then tell me how many orders they have made.",
|
||||
|
|
@ -279,6 +281,7 @@ Args: {
|
|||
|
||||
flow = MagicMock()
|
||||
flow.side_effect = flow_context
|
||||
flow.workspace = "default"
|
||||
|
||||
# Act
|
||||
await agent_processor.on_request(msg, consumer, flow)
|
||||
|
|
@ -313,7 +316,7 @@ Args: {
|
|||
}
|
||||
}
|
||||
|
||||
await agent_processor.on_tools_config(tool_config_with_collection, "v1")
|
||||
await agent_processor.on_tools_config("default", tool_config_with_collection, "v1")
|
||||
|
||||
request = AgentRequest(
|
||||
question="Query the sales data for recent transactions.",
|
||||
|
|
@ -371,6 +374,7 @@ Args: {
|
|||
|
||||
flow = MagicMock()
|
||||
flow.side_effect = flow_context
|
||||
flow.workspace = "default"
|
||||
|
||||
# Act
|
||||
await agent_processor.on_request(msg, consumer, flow)
|
||||
|
|
@ -394,10 +398,10 @@ Args: {
|
|||
async def test_agent_structured_query_tool_argument_validation(self, agent_processor, structured_query_tool_config):
|
||||
"""Test that structured query tool arguments are properly validated"""
|
||||
# Arrange
|
||||
await agent_processor.on_tools_config(structured_query_tool_config, "v1")
|
||||
await agent_processor.on_tools_config("default", structured_query_tool_config, "v1")
|
||||
|
||||
# Check that the tool was registered with correct arguments
|
||||
tools = agent_processor.agent.tools
|
||||
tools = agent_processor.agents["default"].tools
|
||||
assert "structured-query" in tools
|
||||
|
||||
structured_tool = tools["structured-query"]
|
||||
|
|
@ -414,7 +418,7 @@ Args: {
|
|||
async def test_agent_structured_query_json_formatting(self, agent_processor, structured_query_tool_config):
|
||||
"""Test that structured query results are properly formatted for agent consumption"""
|
||||
# Arrange
|
||||
await agent_processor.on_tools_config(structured_query_tool_config, "v1")
|
||||
await agent_processor.on_tools_config("default", structured_query_tool_config, "v1")
|
||||
|
||||
request = AgentRequest(
|
||||
question="Get customer information and format it nicely.",
|
||||
|
|
@ -482,6 +486,7 @@ Args: {
|
|||
|
||||
flow = MagicMock()
|
||||
flow.side_effect = flow_context
|
||||
flow.workspace = "default"
|
||||
|
||||
# Act
|
||||
await agent_processor.on_request(msg, consumer, flow)
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ class TestNLPQueryServiceIntegration:
|
|||
)
|
||||
|
||||
# Set up schemas
|
||||
proc.schemas = sample_schemas
|
||||
proc.schemas = {"default": dict(sample_schemas)}
|
||||
|
||||
# Mock the client method
|
||||
proc.client = MagicMock()
|
||||
|
|
@ -94,6 +94,7 @@ class TestNLPQueryServiceIntegration:
|
|||
|
||||
consumer = MagicMock()
|
||||
flow = MagicMock()
|
||||
flow.workspace = "default"
|
||||
flow_response = AsyncMock()
|
||||
flow.return_value = flow_response
|
||||
|
||||
|
|
@ -173,6 +174,7 @@ class TestNLPQueryServiceIntegration:
|
|||
|
||||
consumer = MagicMock()
|
||||
flow = MagicMock()
|
||||
flow.workspace = "default"
|
||||
flow_response = AsyncMock()
|
||||
flow.return_value = flow_response
|
||||
|
||||
|
|
@ -229,7 +231,7 @@ class TestNLPQueryServiceIntegration:
|
|||
}
|
||||
|
||||
# Act - Update configuration
|
||||
await integration_processor.on_schema_config(new_schema_config, "v2")
|
||||
await integration_processor.on_schema_config("default", new_schema_config, "v2")
|
||||
|
||||
# Arrange - Test query using new schema
|
||||
request = QuestionToStructuredQueryRequest(
|
||||
|
|
@ -243,6 +245,7 @@ class TestNLPQueryServiceIntegration:
|
|||
|
||||
consumer = MagicMock()
|
||||
flow = MagicMock()
|
||||
flow.workspace = "default"
|
||||
flow_response = AsyncMock()
|
||||
flow.return_value = flow_response
|
||||
|
||||
|
|
@ -272,7 +275,7 @@ class TestNLPQueryServiceIntegration:
|
|||
await integration_processor.on_message(msg, consumer, flow)
|
||||
|
||||
# Assert
|
||||
assert "inventory" in integration_processor.schemas
|
||||
assert "inventory" in integration_processor.schemas["default"]
|
||||
response_call = flow_response.send.call_args
|
||||
response = response_call[0][0]
|
||||
assert response.detected_schemas == ["inventory"]
|
||||
|
|
@ -293,6 +296,7 @@ class TestNLPQueryServiceIntegration:
|
|||
|
||||
consumer = MagicMock()
|
||||
flow = MagicMock()
|
||||
flow.workspace = "default"
|
||||
flow_response = AsyncMock()
|
||||
flow.return_value = flow_response
|
||||
|
||||
|
|
@ -334,7 +338,7 @@ class TestNLPQueryServiceIntegration:
|
|||
graphql_generation_template="custom-graphql-generator"
|
||||
)
|
||||
|
||||
custom_processor.schemas = sample_schemas
|
||||
custom_processor.schemas = {"default": dict(sample_schemas)}
|
||||
custom_processor.client = MagicMock()
|
||||
|
||||
request = QuestionToStructuredQueryRequest(
|
||||
|
|
@ -348,6 +352,7 @@ class TestNLPQueryServiceIntegration:
|
|||
|
||||
consumer = MagicMock()
|
||||
flow = MagicMock()
|
||||
flow.workspace = "default"
|
||||
flow_response = AsyncMock()
|
||||
flow.return_value = flow_response
|
||||
|
||||
|
|
@ -394,7 +399,7 @@ class TestNLPQueryServiceIntegration:
|
|||
] + [SchemaField(name=f"field_{j}", type="string") for j in range(5)]
|
||||
)
|
||||
|
||||
integration_processor.schemas.update(large_schema_set)
|
||||
integration_processor.schemas["default"].update(large_schema_set)
|
||||
|
||||
request = QuestionToStructuredQueryRequest(
|
||||
question="Show me data from table_05 and table_12",
|
||||
|
|
@ -407,6 +412,7 @@ class TestNLPQueryServiceIntegration:
|
|||
|
||||
consumer = MagicMock()
|
||||
flow = MagicMock()
|
||||
flow.workspace = "default"
|
||||
flow_response = AsyncMock()
|
||||
flow.return_value = flow_response
|
||||
|
||||
|
|
@ -462,6 +468,7 @@ class TestNLPQueryServiceIntegration:
|
|||
msg.properties.return_value = {"id": f"concurrent-test-{i}"}
|
||||
|
||||
flow = MagicMock()
|
||||
flow.workspace = "default"
|
||||
flow_response = AsyncMock()
|
||||
flow.return_value = flow_response
|
||||
|
||||
|
|
@ -532,6 +539,7 @@ class TestNLPQueryServiceIntegration:
|
|||
|
||||
consumer = MagicMock()
|
||||
flow = MagicMock()
|
||||
flow.workspace = "default"
|
||||
flow_response = AsyncMock()
|
||||
flow.return_value = flow_response
|
||||
|
||||
|
|
|
|||
|
|
@ -185,6 +185,7 @@ class TestObjectExtractionServiceIntegration:
|
|||
return AsyncMock()
|
||||
|
||||
context.side_effect = context_router
|
||||
context.workspace = "default"
|
||||
return context
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -197,20 +198,21 @@ class TestObjectExtractionServiceIntegration:
|
|||
processor.on_schema_config = Processor.on_schema_config.__get__(processor, Processor)
|
||||
|
||||
# Act
|
||||
await processor.on_schema_config(integration_config, version=1)
|
||||
await processor.on_schema_config("default", integration_config, version=1)
|
||||
|
||||
# Assert
|
||||
assert len(processor.schemas) == 2
|
||||
assert "customer_records" in processor.schemas
|
||||
assert "product_catalog" in processor.schemas
|
||||
|
||||
ws_schemas = processor.schemas["default"]
|
||||
assert len(ws_schemas) == 2
|
||||
assert "customer_records" in ws_schemas
|
||||
assert "product_catalog" in ws_schemas
|
||||
|
||||
# Verify customer schema
|
||||
customer_schema = processor.schemas["customer_records"]
|
||||
customer_schema = ws_schemas["customer_records"]
|
||||
assert customer_schema.name == "customer_records"
|
||||
assert len(customer_schema.fields) == 4
|
||||
|
||||
|
||||
# Verify product schema
|
||||
product_schema = processor.schemas["product_catalog"]
|
||||
product_schema = ws_schemas["product_catalog"]
|
||||
assert product_schema.name == "product_catalog"
|
||||
assert len(product_schema.fields) == 4
|
||||
|
||||
|
|
@ -237,7 +239,7 @@ class TestObjectExtractionServiceIntegration:
|
|||
processor.convert_values_to_strings = convert_values_to_strings
|
||||
|
||||
# Load configuration
|
||||
await processor.on_schema_config(integration_config, version=1)
|
||||
await processor.on_schema_config("default", integration_config, version=1)
|
||||
|
||||
# Create realistic customer data chunk
|
||||
metadata = Metadata(
|
||||
|
|
@ -304,7 +306,7 @@ class TestObjectExtractionServiceIntegration:
|
|||
processor.convert_values_to_strings = convert_values_to_strings
|
||||
|
||||
# Load configuration
|
||||
await processor.on_schema_config(integration_config, version=1)
|
||||
await processor.on_schema_config("default", integration_config, version=1)
|
||||
|
||||
# Create realistic product data chunk
|
||||
metadata = Metadata(
|
||||
|
|
@ -368,7 +370,7 @@ class TestObjectExtractionServiceIntegration:
|
|||
processor.convert_values_to_strings = convert_values_to_strings
|
||||
|
||||
# Load configuration
|
||||
await processor.on_schema_config(integration_config, version=1)
|
||||
await processor.on_schema_config("default", integration_config, version=1)
|
||||
|
||||
# Create multiple test chunks
|
||||
chunks_data = [
|
||||
|
|
@ -431,19 +433,21 @@ class TestObjectExtractionServiceIntegration:
|
|||
"customer_records": integration_config["schema"]["customer_records"]
|
||||
}
|
||||
}
|
||||
await processor.on_schema_config(initial_config, version=1)
|
||||
|
||||
assert len(processor.schemas) == 1
|
||||
assert "customer_records" in processor.schemas
|
||||
assert "product_catalog" not in processor.schemas
|
||||
|
||||
await processor.on_schema_config("default", initial_config, version=1)
|
||||
|
||||
ws_schemas = processor.schemas["default"]
|
||||
assert len(ws_schemas) == 1
|
||||
assert "customer_records" in ws_schemas
|
||||
assert "product_catalog" not in ws_schemas
|
||||
|
||||
# Act - Reload with full configuration
|
||||
await processor.on_schema_config(integration_config, version=2)
|
||||
|
||||
await processor.on_schema_config("default", integration_config, version=2)
|
||||
|
||||
# Assert
|
||||
assert len(processor.schemas) == 2
|
||||
assert "customer_records" in processor.schemas
|
||||
assert "product_catalog" in processor.schemas
|
||||
ws_schemas = processor.schemas["default"]
|
||||
assert len(ws_schemas) == 2
|
||||
assert "customer_records" in ws_schemas
|
||||
assert "product_catalog" in ws_schemas
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_resilience_integration(self, integration_config):
|
||||
|
|
@ -474,10 +478,11 @@ class TestObjectExtractionServiceIntegration:
|
|||
return AsyncMock()
|
||||
|
||||
failing_flow.side_effect = failing_context_router
|
||||
failing_flow.workspace = "default"
|
||||
processor.flow = failing_flow
|
||||
|
||||
# Load configuration
|
||||
await processor.on_schema_config(integration_config, version=1)
|
||||
await processor.on_schema_config("default", integration_config, version=1)
|
||||
|
||||
# Create test chunk
|
||||
metadata = Metadata(id="error-test", user="test", collection="test")
|
||||
|
|
@ -510,7 +515,7 @@ class TestObjectExtractionServiceIntegration:
|
|||
processor.convert_values_to_strings = convert_values_to_strings
|
||||
|
||||
# Load configuration
|
||||
await processor.on_schema_config(integration_config, version=1)
|
||||
await processor.on_schema_config("default", integration_config, version=1)
|
||||
|
||||
# Create chunk with rich metadata
|
||||
original_metadata = Metadata(
|
||||
|
|
|
|||
|
|
@ -87,6 +87,7 @@ class TestPromptStreaming:
|
|||
return AsyncMock()
|
||||
|
||||
context.side_effect = context_router
|
||||
context.workspace = "default"
|
||||
return context
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -109,7 +110,7 @@ class TestPromptStreaming:
|
|||
def prompt_processor_streaming(self, mock_prompt_manager):
|
||||
"""Create Prompt processor with streaming support"""
|
||||
processor = MagicMock()
|
||||
processor.manager = mock_prompt_manager
|
||||
processor.managers = {"default": mock_prompt_manager}
|
||||
processor.config_key = "prompt"
|
||||
|
||||
# Bind the actual on_request method
|
||||
|
|
@ -248,6 +249,7 @@ class TestPromptStreaming:
|
|||
return AsyncMock()
|
||||
|
||||
context.side_effect = context_router
|
||||
context.workspace = "default"
|
||||
|
||||
request = PromptRequest(
|
||||
id="test_prompt",
|
||||
|
|
@ -341,6 +343,7 @@ class TestPromptStreaming:
|
|||
return AsyncMock()
|
||||
|
||||
context.side_effect = context_router
|
||||
context.workspace = "default"
|
||||
|
||||
request = PromptRequest(
|
||||
id="test_prompt",
|
||||
|
|
|
|||
|
|
@ -14,6 +14,17 @@ from trustgraph.storage.rows.cassandra.write import Processor
|
|||
from trustgraph.schema import ExtractedObject, Metadata, RowSchema, Field
|
||||
|
||||
|
||||
|
||||
|
||||
class _MockFlowDefault:
|
||||
"""Mock Flow with default workspace for testing."""
|
||||
workspace = "default"
|
||||
name = "default"
|
||||
id = "test-processor"
|
||||
|
||||
|
||||
mock_flow_default = _MockFlowDefault()
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestRowsCassandraIntegration:
|
||||
"""Integration tests for Cassandra row storage with unified table"""
|
||||
|
|
@ -125,8 +136,8 @@ class TestRowsCassandraIntegration:
|
|||
}
|
||||
}
|
||||
|
||||
await processor.on_schema_config(config, version=1)
|
||||
assert "customer_records" in processor.schemas
|
||||
await processor.on_schema_config("default", config, version=1)
|
||||
assert "customer_records" in processor.schemas["default"]
|
||||
|
||||
# Step 2: Process an ExtractedObject
|
||||
test_obj = ExtractedObject(
|
||||
|
|
@ -149,7 +160,7 @@ class TestRowsCassandraIntegration:
|
|||
msg = MagicMock()
|
||||
msg.value.return_value = test_obj
|
||||
|
||||
await processor.on_object(msg, None, None)
|
||||
await processor.on_object(msg, None, mock_flow_default)
|
||||
|
||||
# Verify Cassandra interactions
|
||||
assert mock_cluster.connect.called
|
||||
|
|
@ -209,8 +220,8 @@ class TestRowsCassandraIntegration:
|
|||
}
|
||||
}
|
||||
|
||||
await processor.on_schema_config(config, version=1)
|
||||
assert len(processor.schemas) == 2
|
||||
await processor.on_schema_config("default", config, version=1)
|
||||
assert len(processor.schemas["default"]) == 2
|
||||
|
||||
# Process objects for different schemas
|
||||
product_obj = ExtractedObject(
|
||||
|
|
@ -233,7 +244,7 @@ class TestRowsCassandraIntegration:
|
|||
for obj in [product_obj, order_obj]:
|
||||
msg = MagicMock()
|
||||
msg.value.return_value = obj
|
||||
await processor.on_object(msg, None, None)
|
||||
await processor.on_object(msg, None, mock_flow_default)
|
||||
|
||||
# All data goes into the same unified rows table
|
||||
table_calls = [call for call in mock_session.execute.call_args_list
|
||||
|
|
@ -256,15 +267,17 @@ class TestRowsCassandraIntegration:
|
|||
|
||||
with patch('trustgraph.storage.rows.cassandra.write.Cluster', return_value=mock_cluster):
|
||||
# Schema with multiple indexed fields
|
||||
processor.schemas["indexed_data"] = RowSchema(
|
||||
name="indexed_data",
|
||||
fields=[
|
||||
Field(name="id", type="string", size=50, primary=True),
|
||||
Field(name="category", type="string", size=50, indexed=True),
|
||||
Field(name="status", type="string", size=50, indexed=True),
|
||||
Field(name="description", type="string", size=200) # Not indexed
|
||||
]
|
||||
)
|
||||
processor.schemas["default"] = {
|
||||
"indexed_data": RowSchema(
|
||||
name="indexed_data",
|
||||
fields=[
|
||||
Field(name="id", type="string", size=50, primary=True),
|
||||
Field(name="category", type="string", size=50, indexed=True),
|
||||
Field(name="status", type="string", size=50, indexed=True),
|
||||
Field(name="description", type="string", size=200) # Not indexed
|
||||
]
|
||||
)
|
||||
}
|
||||
|
||||
test_obj = ExtractedObject(
|
||||
metadata=Metadata(id="t1", user="test", collection="test"),
|
||||
|
|
@ -282,7 +295,7 @@ class TestRowsCassandraIntegration:
|
|||
msg = MagicMock()
|
||||
msg.value.return_value = test_obj
|
||||
|
||||
await processor.on_object(msg, None, None)
|
||||
await processor.on_object(msg, None, mock_flow_default)
|
||||
|
||||
# Should have 3 data inserts (one per indexed field: id, category, status)
|
||||
rows_insert_calls = [call for call in mock_session.execute.call_args_list
|
||||
|
|
@ -342,7 +355,7 @@ class TestRowsCassandraIntegration:
|
|||
}
|
||||
}
|
||||
|
||||
await processor.on_schema_config(config, version=1)
|
||||
await processor.on_schema_config("default", config, version=1)
|
||||
|
||||
# Process batch object with multiple values
|
||||
batch_obj = ExtractedObject(
|
||||
|
|
@ -376,7 +389,7 @@ class TestRowsCassandraIntegration:
|
|||
msg = MagicMock()
|
||||
msg.value.return_value = batch_obj
|
||||
|
||||
await processor.on_object(msg, None, None)
|
||||
await processor.on_object(msg, None, mock_flow_default)
|
||||
|
||||
# Verify unified table creation
|
||||
table_calls = [call for call in mock_session.execute.call_args_list
|
||||
|
|
@ -396,10 +409,12 @@ class TestRowsCassandraIntegration:
|
|||
processor, mock_cluster, mock_session = processor_with_mocks
|
||||
|
||||
with patch('trustgraph.storage.rows.cassandra.write.Cluster', return_value=mock_cluster):
|
||||
processor.schemas["empty_test"] = RowSchema(
|
||||
name="empty_test",
|
||||
fields=[Field(name="id", type="string", size=50, primary=True)]
|
||||
)
|
||||
processor.schemas["default"] = {
|
||||
"empty_test": RowSchema(
|
||||
name="empty_test",
|
||||
fields=[Field(name="id", type="string", size=50, primary=True)]
|
||||
)
|
||||
}
|
||||
|
||||
# Process empty batch object
|
||||
empty_obj = ExtractedObject(
|
||||
|
|
@ -413,7 +428,7 @@ class TestRowsCassandraIntegration:
|
|||
msg = MagicMock()
|
||||
msg.value.return_value = empty_obj
|
||||
|
||||
await processor.on_object(msg, None, None)
|
||||
await processor.on_object(msg, None, mock_flow_default)
|
||||
|
||||
# Should not create any data insert statements for empty batch
|
||||
# (partition registration may still happen)
|
||||
|
|
@ -428,14 +443,16 @@ class TestRowsCassandraIntegration:
|
|||
processor, mock_cluster, mock_session = processor_with_mocks
|
||||
|
||||
with patch('trustgraph.storage.rows.cassandra.write.Cluster', return_value=mock_cluster):
|
||||
processor.schemas["map_test"] = RowSchema(
|
||||
name="map_test",
|
||||
fields=[
|
||||
Field(name="id", type="string", size=50, primary=True),
|
||||
Field(name="name", type="string", size=100),
|
||||
Field(name="count", type="integer", size=0)
|
||||
]
|
||||
)
|
||||
processor.schemas["default"] = {
|
||||
"map_test": RowSchema(
|
||||
name="map_test",
|
||||
fields=[
|
||||
Field(name="id", type="string", size=50, primary=True),
|
||||
Field(name="name", type="string", size=100),
|
||||
Field(name="count", type="integer", size=0)
|
||||
]
|
||||
)
|
||||
}
|
||||
|
||||
test_obj = ExtractedObject(
|
||||
metadata=Metadata(id="t1", user="test", collection="test"),
|
||||
|
|
@ -448,7 +465,7 @@ class TestRowsCassandraIntegration:
|
|||
msg = MagicMock()
|
||||
msg.value.return_value = test_obj
|
||||
|
||||
await processor.on_object(msg, None, None)
|
||||
await processor.on_object(msg, None, mock_flow_default)
|
||||
|
||||
# Verify insert uses map for data
|
||||
rows_insert_calls = [call for call in mock_session.execute.call_args_list
|
||||
|
|
@ -473,13 +490,15 @@ class TestRowsCassandraIntegration:
|
|||
processor, mock_cluster, mock_session = processor_with_mocks
|
||||
|
||||
with patch('trustgraph.storage.rows.cassandra.write.Cluster', return_value=mock_cluster):
|
||||
processor.schemas["partition_test"] = RowSchema(
|
||||
name="partition_test",
|
||||
fields=[
|
||||
Field(name="id", type="string", size=50, primary=True),
|
||||
Field(name="category", type="string", size=50, indexed=True)
|
||||
]
|
||||
)
|
||||
processor.schemas["default"] = {
|
||||
"partition_test": RowSchema(
|
||||
name="partition_test",
|
||||
fields=[
|
||||
Field(name="id", type="string", size=50, primary=True),
|
||||
Field(name="category", type="string", size=50, indexed=True)
|
||||
]
|
||||
)
|
||||
}
|
||||
|
||||
test_obj = ExtractedObject(
|
||||
metadata=Metadata(id="t1", user="test", collection="my_collection"),
|
||||
|
|
@ -492,7 +511,7 @@ class TestRowsCassandraIntegration:
|
|||
msg = MagicMock()
|
||||
msg.value.return_value = test_obj
|
||||
|
||||
await processor.on_object(msg, None, None)
|
||||
await processor.on_object(msg, None, mock_flow_default)
|
||||
|
||||
# Verify partition registration
|
||||
partition_inserts = [call for call in mock_session.execute.call_args_list
|
||||
|
|
|
|||
|
|
@ -154,7 +154,7 @@ class TestObjectsGraphQLQueryIntegration:
|
|||
async def test_schema_configuration_and_generation(self, processor, sample_schema_config):
|
||||
"""Test schema configuration loading and GraphQL schema generation"""
|
||||
# Load schema configuration
|
||||
await processor.on_schema_config(sample_schema_config, version=1)
|
||||
await processor.on_schema_config("default", sample_schema_config, version=1)
|
||||
|
||||
# Verify schemas were loaded
|
||||
assert len(processor.schemas) == 2
|
||||
|
|
@ -181,7 +181,7 @@ class TestObjectsGraphQLQueryIntegration:
|
|||
async def test_cassandra_connection_and_table_creation(self, processor, sample_schema_config):
|
||||
"""Test Cassandra connection and dynamic table creation"""
|
||||
# Load schema configuration
|
||||
await processor.on_schema_config(sample_schema_config, version=1)
|
||||
await processor.on_schema_config("default", sample_schema_config, version=1)
|
||||
|
||||
# Connect to Cassandra
|
||||
processor.connect_cassandra()
|
||||
|
|
@ -218,7 +218,7 @@ class TestObjectsGraphQLQueryIntegration:
|
|||
async def test_data_insertion_and_graphql_query(self, processor, sample_schema_config):
|
||||
"""Test inserting data and querying via GraphQL"""
|
||||
# Load schema and connect
|
||||
await processor.on_schema_config(sample_schema_config, version=1)
|
||||
await processor.on_schema_config("default", sample_schema_config, version=1)
|
||||
processor.connect_cassandra()
|
||||
|
||||
# Setup test data
|
||||
|
|
@ -292,7 +292,7 @@ class TestObjectsGraphQLQueryIntegration:
|
|||
async def test_graphql_query_with_filters(self, processor, sample_schema_config):
|
||||
"""Test GraphQL queries with filtering on indexed fields"""
|
||||
# Setup (reuse previous setup)
|
||||
await processor.on_schema_config(sample_schema_config, version=1)
|
||||
await processor.on_schema_config("default", sample_schema_config, version=1)
|
||||
processor.connect_cassandra()
|
||||
|
||||
keyspace = "test_user"
|
||||
|
|
@ -353,7 +353,7 @@ class TestObjectsGraphQLQueryIntegration:
|
|||
async def test_graphql_error_handling(self, processor, sample_schema_config):
|
||||
"""Test GraphQL error handling for invalid queries"""
|
||||
# Setup
|
||||
await processor.on_schema_config(sample_schema_config, version=1)
|
||||
await processor.on_schema_config("default", sample_schema_config, version=1)
|
||||
|
||||
# Test invalid field query
|
||||
invalid_query = '''
|
||||
|
|
@ -386,7 +386,7 @@ class TestObjectsGraphQLQueryIntegration:
|
|||
async def test_message_processing_integration(self, processor, sample_schema_config):
|
||||
"""Test full message processing workflow"""
|
||||
# Setup
|
||||
await processor.on_schema_config(sample_schema_config, version=1)
|
||||
await processor.on_schema_config("default", sample_schema_config, version=1)
|
||||
processor.connect_cassandra()
|
||||
|
||||
# Create mock message
|
||||
|
|
@ -432,7 +432,7 @@ class TestObjectsGraphQLQueryIntegration:
|
|||
async def test_concurrent_queries(self, processor, sample_schema_config):
|
||||
"""Test handling multiple concurrent GraphQL queries"""
|
||||
# Setup
|
||||
await processor.on_schema_config(sample_schema_config, version=1)
|
||||
await processor.on_schema_config("default", sample_schema_config, version=1)
|
||||
processor.connect_cassandra()
|
||||
|
||||
# Create multiple query tasks
|
||||
|
|
@ -476,7 +476,7 @@ class TestObjectsGraphQLQueryIntegration:
|
|||
}
|
||||
}
|
||||
|
||||
await processor.on_schema_config(initial_config, version=1)
|
||||
await processor.on_schema_config("default", initial_config, version=1)
|
||||
assert len(processor.schemas) == 1
|
||||
assert "simple" in processor.schemas
|
||||
|
||||
|
|
@ -500,7 +500,7 @@ class TestObjectsGraphQLQueryIntegration:
|
|||
}
|
||||
}
|
||||
|
||||
await processor.on_schema_config(updated_config, version=2)
|
||||
await processor.on_schema_config("default", updated_config, version=2)
|
||||
|
||||
# Verify updated schemas
|
||||
assert len(processor.schemas) == 2
|
||||
|
|
@ -518,7 +518,7 @@ class TestObjectsGraphQLQueryIntegration:
|
|||
async def test_large_result_set_handling(self, processor, sample_schema_config):
|
||||
"""Test handling of large query result sets"""
|
||||
# Setup
|
||||
await processor.on_schema_config(sample_schema_config, version=1)
|
||||
await processor.on_schema_config("default", sample_schema_config, version=1)
|
||||
processor.connect_cassandra()
|
||||
|
||||
keyspace = "large_test_user"
|
||||
|
|
@ -601,7 +601,7 @@ class TestObjectsGraphQLQueryPerformance:
|
|||
}
|
||||
}
|
||||
|
||||
await processor.on_schema_config(schema_config, version=1)
|
||||
await processor.on_schema_config("default", schema_config, version=1)
|
||||
|
||||
# Measure query execution time
|
||||
start_time = time.time()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue