Basic multitenant support (#583)

* Tech spec

* Address multi-tenant queue option problems in CLI

* Modified collection service to use config

* Changed storage management to use the config service definition
This commit is contained in:
cybermaggedon 2025-12-05 21:45:30 +00:00 committed by GitHub
parent 789d9713a0
commit 7d07f802a8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
28 changed files with 1416 additions and 1731 deletions

View file

@ -145,7 +145,7 @@ class TestResolveCassandraConfig:
def test_default_configuration(self):
"""Test resolution with no parameters or environment variables."""
with patch.dict(os.environ, {}, clear=True):
hosts, username, password = resolve_cassandra_config()
hosts, username, password, keyspace = resolve_cassandra_config()
assert hosts == ['cassandra']
assert username is None
@ -160,7 +160,7 @@ class TestResolveCassandraConfig:
}
with patch.dict(os.environ, env_vars, clear=True):
hosts, username, password = resolve_cassandra_config()
hosts, username, password, keyspace = resolve_cassandra_config()
assert hosts == ['env1', 'env2', 'env3']
assert username == 'env-user'
@ -175,7 +175,7 @@ class TestResolveCassandraConfig:
}
with patch.dict(os.environ, env_vars, clear=True):
hosts, username, password = resolve_cassandra_config(
hosts, username, password, keyspace = resolve_cassandra_config(
host='explicit-host',
username='explicit-user',
password='explicit-pass'
@ -188,19 +188,19 @@ class TestResolveCassandraConfig:
def test_host_list_parsing(self):
"""Test different host list formats."""
# Single host
hosts, _, _ = resolve_cassandra_config(host='single-host')
hosts, _, _, _ = resolve_cassandra_config(host='single-host')
assert hosts == ['single-host']
# Multiple hosts with spaces
hosts, _, _ = resolve_cassandra_config(host='host1, host2 ,host3')
hosts, _, _, _ = resolve_cassandra_config(host='host1, host2 ,host3')
assert hosts == ['host1', 'host2', 'host3']
# Empty elements filtered out
hosts, _, _ = resolve_cassandra_config(host='host1,,host2,')
hosts, _, _, _ = resolve_cassandra_config(host='host1,,host2,')
assert hosts == ['host1', 'host2']
# Already a list
hosts, _, _ = resolve_cassandra_config(host=['list-host1', 'list-host2'])
hosts, _, _, _ = resolve_cassandra_config(host=['list-host1', 'list-host2'])
assert hosts == ['list-host1', 'list-host2']
def test_args_object_resolution(self):
@ -212,7 +212,7 @@ class TestResolveCassandraConfig:
cassandra_password = 'args-pass'
args = MockArgs()
hosts, username, password = resolve_cassandra_config(args)
hosts, username, password, keyspace = resolve_cassandra_config(args)
assert hosts == ['args-host1', 'args-host2']
assert username == 'args-user'
@ -233,7 +233,7 @@ class TestResolveCassandraConfig:
with patch.dict(os.environ, env_vars, clear=True):
args = PartialArgs()
hosts, username, password = resolve_cassandra_config(args)
hosts, username, password, keyspace = resolve_cassandra_config(args)
assert hosts == ['args-host'] # From args
assert username == 'env-user' # From env
@ -251,7 +251,7 @@ class TestGetCassandraConfigFromParams:
'cassandra_password': 'new-pass'
}
hosts, username, password = get_cassandra_config_from_params(params)
hosts, username, password, keyspace = get_cassandra_config_from_params(params)
assert hosts == ['new-host1', 'new-host2']
assert username == 'new-user'
@ -265,7 +265,7 @@ class TestGetCassandraConfigFromParams:
'graph_password': 'old-pass'
}
hosts, username, password = get_cassandra_config_from_params(params)
hosts, username, password, keyspace = get_cassandra_config_from_params(params)
# Should use defaults since graph_* params are not recognized
assert hosts == ['cassandra'] # Default
@ -280,7 +280,7 @@ class TestGetCassandraConfigFromParams:
'cassandra_password': 'compat-pass'
}
hosts, username, password = get_cassandra_config_from_params(params)
hosts, username, password, keyspace = get_cassandra_config_from_params(params)
assert hosts == ['compat-host']
assert username is None # cassandra_user is not recognized
@ -298,7 +298,7 @@ class TestGetCassandraConfigFromParams:
'graph_password': 'old-pass'
}
hosts, username, password = get_cassandra_config_from_params(params)
hosts, username, password, keyspace = get_cassandra_config_from_params(params)
assert hosts == ['new-host'] # Only cassandra_* params work
assert username == 'new-user' # Only cassandra_* params work
@ -314,7 +314,7 @@ class TestGetCassandraConfigFromParams:
with patch.dict(os.environ, env_vars, clear=True):
params = {}
hosts, username, password = get_cassandra_config_from_params(params)
hosts, username, password, keyspace = get_cassandra_config_from_params(params)
assert hosts == ['fallback-host1', 'fallback-host2']
assert username == 'fallback-user'
@ -334,7 +334,7 @@ class TestConfigurationPriority:
with patch.dict(os.environ, env_vars, clear=True):
# CLI args should override everything
hosts, username, password = resolve_cassandra_config(
hosts, username, password, keyspace = resolve_cassandra_config(
host='cli-host',
username='cli-user',
password='cli-pass'
@ -354,7 +354,7 @@ class TestConfigurationPriority:
with patch.dict(os.environ, env_vars, clear=True):
# Only provide host via CLI
hosts, username, password = resolve_cassandra_config(
hosts, username, password, keyspace = resolve_cassandra_config(
host='cli-host'
# username and password not provided
)
@ -366,7 +366,7 @@ class TestConfigurationPriority:
def test_no_config_defaults(self):
"""Test that defaults are used when no configuration is provided."""
with patch.dict(os.environ, {}, clear=True):
hosts, username, password = resolve_cassandra_config()
hosts, username, password, keyspace = resolve_cassandra_config()
assert hosts == ['cassandra'] # Default
assert username is None # Default
@ -378,17 +378,17 @@ class TestEdgeCases:
def test_empty_host_string(self):
"""Test handling of empty host string falls back to default."""
hosts, _, _ = resolve_cassandra_config(host='')
hosts, _, _, _ = resolve_cassandra_config(host='')
assert hosts == ['cassandra'] # Falls back to default
def test_whitespace_only_host(self):
"""Test handling of whitespace-only host string."""
hosts, _, _ = resolve_cassandra_config(host=' ')
hosts, _, _, _ = resolve_cassandra_config(host=' ')
assert hosts == [] # Empty after stripping whitespace
def test_none_values_preserved(self):
"""Test that None values are preserved correctly."""
hosts, username, password = resolve_cassandra_config(
hosts, username, password, keyspace = resolve_cassandra_config(
host=None,
username=None,
password=None
@ -401,7 +401,7 @@ class TestEdgeCases:
def test_mixed_none_and_values(self):
"""Test mixing None and actual values."""
hosts, username, password = resolve_cassandra_config(
hosts, username, password, keyspace = resolve_cassandra_config(
host='mixed-host',
username=None,
password='mixed-pass'

View file

@ -15,11 +15,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
"""Test Qdrant document embeddings storage functionality"""
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_processor_initialization_basic(self, mock_base_init, mock_qdrant_client):
async def test_processor_initialization_basic(self, mock_qdrant_client):
"""Test basic Qdrant processor initialization"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_client.return_value = mock_qdrant_instance
@ -34,9 +32,6 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
processor = Processor(**config)
# Assert
# Verify base class initialization was called
mock_base_init.assert_called_once()
# Verify QdrantClient was created with correct parameters
mock_qdrant_client.assert_called_once_with(url='http://localhost:6333', api_key='test-api-key')
@ -45,11 +40,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
assert processor.qdrant == mock_qdrant_instance
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_processor_initialization_with_defaults(self, mock_base_init, mock_qdrant_client):
async def test_processor_initialization_with_defaults(self, mock_qdrant_client):
"""Test processor initialization with default values"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_client.return_value = mock_qdrant_instance
@ -68,11 +61,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_store_document_embeddings_basic(self, mock_base_init, mock_uuid, mock_qdrant_client):
async def test_store_document_embeddings_basic(self, mock_uuid, mock_qdrant_client):
"""Test storing document embeddings with basic message"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True # Collection already exists
mock_qdrant_client.return_value = mock_qdrant_instance
@ -121,11 +112,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_store_document_embeddings_multiple_chunks(self, mock_base_init, mock_uuid, mock_qdrant_client):
async def test_store_document_embeddings_multiple_chunks(self, mock_uuid, mock_qdrant_client):
"""Test storing document embeddings with multiple chunks"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True
mock_qdrant_client.return_value = mock_qdrant_instance
@ -180,11 +169,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_store_document_embeddings_multiple_vectors_per_chunk(self, mock_base_init, mock_uuid, mock_qdrant_client):
async def test_store_document_embeddings_multiple_vectors_per_chunk(self, mock_uuid, mock_qdrant_client):
"""Test storing document embeddings with multiple vectors per chunk"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True
mock_qdrant_client.return_value = mock_qdrant_instance
@ -237,11 +224,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
assert point.payload['doc'] == 'multi-vector document chunk'
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_store_document_embeddings_empty_chunk(self, mock_base_init, mock_qdrant_client):
async def test_store_document_embeddings_empty_chunk(self, mock_qdrant_client):
"""Test storing document embeddings skips empty chunks"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True # Collection exists
mock_qdrant_client.return_value = mock_qdrant_instance
@ -277,11 +262,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_collection_creation_when_not_exists(self, mock_base_init, mock_uuid, mock_qdrant_client):
async def test_collection_creation_when_not_exists(self, mock_uuid, mock_qdrant_client):
"""Test that writing to non-existent collection creates it lazily"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = False # Collection doesn't exist
mock_qdrant_client.return_value = mock_qdrant_instance
@ -326,11 +309,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_collection_creation_exception(self, mock_base_init, mock_uuid, mock_qdrant_client):
async def test_collection_creation_exception(self, mock_uuid, mock_qdrant_client):
"""Test that collection creation errors are propagated"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = False # Collection doesn't exist
# Simulate creation failure
@ -364,12 +345,10 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
await processor.store_document_embeddings(mock_message)
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
async def test_collection_validation_on_write(self, mock_uuid, mock_base_init, mock_qdrant_client):
async def test_collection_validation_on_write(self, mock_uuid, mock_qdrant_client):
"""Test collection validation checks collection exists before writing"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True
mock_qdrant_client.return_value = mock_qdrant_instance
@ -428,11 +407,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_different_dimensions_different_collections(self, mock_base_init, mock_uuid, mock_qdrant_client):
async def test_different_dimensions_different_collections(self, mock_uuid, mock_qdrant_client):
"""Test that different vector dimensions create different collections"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True
mock_qdrant_client.return_value = mock_qdrant_instance
@ -482,11 +459,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
assert upsert_calls[1][1]['collection_name'] == 'd_dim_user_dim_collection_3'
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_add_args_calls_parent(self, mock_base_init, mock_qdrant_client):
async def test_add_args_calls_parent(self, mock_qdrant_client):
"""Test that add_args() calls parent add_args method"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_client.return_value = MagicMock()
mock_parser = MagicMock()
@ -502,11 +477,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_utf8_decoding_handling(self, mock_base_init, mock_uuid, mock_qdrant_client):
async def test_utf8_decoding_handling(self, mock_uuid, mock_qdrant_client):
"""Test proper UTF-8 decoding of chunk text"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True
mock_qdrant_client.return_value = mock_qdrant_instance
@ -546,11 +519,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
assert point.payload['doc'] == 'UTF-8 text with special chars: café, naïve, résumé'
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_chunk_decode_exception_handling(self, mock_base_init, mock_qdrant_client):
async def test_chunk_decode_exception_handling(self, mock_qdrant_client):
"""Test handling of chunk decode exceptions"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_client.return_value = mock_qdrant_instance

View file

@ -15,11 +15,9 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
"""Test Qdrant graph embeddings storage functionality"""
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
async def test_processor_initialization_basic(self, mock_base_init, mock_qdrant_client):
async def test_processor_initialization_basic(self, mock_qdrant_client):
"""Test basic Qdrant processor initialization"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_client.return_value = mock_qdrant_instance
@ -34,9 +32,6 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
processor = Processor(**config)
# Assert
# Verify base class initialization was called
mock_base_init.assert_called_once()
# Verify QdrantClient was created with correct parameters
mock_qdrant_client.assert_called_once_with(url='http://localhost:6333', api_key='test-api-key')
@ -46,11 +41,9 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid')
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
async def test_store_graph_embeddings_basic(self, mock_base_init, mock_uuid, mock_qdrant_client):
async def test_store_graph_embeddings_basic(self, mock_uuid, mock_qdrant_client):
"""Test storing graph embeddings with basic message"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True # Collection already exists
mock_qdrant_client.return_value = mock_qdrant_instance
@ -98,11 +91,9 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid')
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
async def test_store_graph_embeddings_multiple_entities(self, mock_base_init, mock_uuid, mock_qdrant_client):
async def test_store_graph_embeddings_multiple_entities(self, mock_uuid, mock_qdrant_client):
"""Test storing graph embeddings with multiple entities"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True
mock_qdrant_client.return_value = mock_qdrant_instance
@ -156,11 +147,9 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid')
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
async def test_store_graph_embeddings_multiple_vectors_per_entity(self, mock_base_init, mock_uuid, mock_qdrant_client):
async def test_store_graph_embeddings_multiple_vectors_per_entity(self, mock_uuid, mock_qdrant_client):
"""Test storing graph embeddings with multiple vectors per entity"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True
mock_qdrant_client.return_value = mock_qdrant_instance
@ -212,11 +201,9 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
assert point.payload['entity'] == 'multi_vector_entity'
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
async def test_store_graph_embeddings_empty_entity_value(self, mock_base_init, mock_qdrant_client):
async def test_store_graph_embeddings_empty_entity_value(self, mock_qdrant_client):
"""Test storing graph embeddings skips empty entity values"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_client.return_value = mock_qdrant_instance
@ -253,11 +240,9 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
mock_qdrant_instance.collection_exists.assert_not_called()
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
async def test_processor_initialization_with_defaults(self, mock_base_init, mock_qdrant_client):
async def test_processor_initialization_with_defaults(self, mock_qdrant_client):
"""Test processor initialization with default values"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_instance = MagicMock()
mock_qdrant_client.return_value = mock_qdrant_instance
@ -275,11 +260,9 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
mock_qdrant_client.assert_called_once_with(url='http://localhost:6333', api_key=None)
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
async def test_add_args_calls_parent(self, mock_base_init, mock_qdrant_client):
async def test_add_args_calls_parent(self, mock_qdrant_client):
"""Test that add_args() calls parent add_args method"""
# Arrange
mock_base_init.return_value = None
mock_qdrant_client.return_value = MagicMock()
mock_parser = MagicMock()