547 lines
19 KiB
Python
547 lines
19 KiB
Python
"""
|
|
Tests for SQLite database manager operations.
|
|
"""
|
|
|
|
import pytest
|
|
import tempfile
|
|
import json
|
|
from pathlib import Path
|
|
from unittest.mock import patch, MagicMock
|
|
import numpy as np
|
|
|
|
from app.core.sqlite_manager import SQLiteManager, generate_slug, get_sqlite_manager
|
|
from app.core.kb_types import KBType
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSQLiteManager:
|
|
"""Test SQLite manager basic operations."""
|
|
|
|
def test_generate_slug(self):
|
|
"""Test slug generation function."""
|
|
test_cases = [
|
|
("My Agent OS", "my-agent-os"),
|
|
("My Code Base!", "my-code-base"),
|
|
("Test_KB 123", "test-kb-123"),
|
|
("Hello World", "hello-world"),
|
|
("Multiple Spaces", "multiple-spaces"),
|
|
("Special@#$%Chars", "specialchars"),
|
|
("", ""),
|
|
]
|
|
|
|
for input_name, expected_slug in test_cases:
|
|
result = generate_slug(input_name)
|
|
assert result == expected_slug, f"Failed for input: {input_name}"
|
|
|
|
def test_sqlite_manager_initialization(self):
|
|
"""Test SQLite manager initialization."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "test.db"
|
|
manager = SQLiteManager(str(db_path))
|
|
|
|
assert manager.db_path == str(db_path)
|
|
assert Path(db_path).exists()
|
|
|
|
def test_create_collection(self, clean_db):
|
|
"""Test creating a knowledge base collection."""
|
|
result = clean_db.create_collection(
|
|
name="test_collection",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test collection",
|
|
tags=["test", "example"]
|
|
)
|
|
|
|
assert result["name"] == "test_collection"
|
|
assert result["kb_type"] == "generic"
|
|
assert result["description"] == "Test collection"
|
|
assert "id" in result
|
|
assert "slug" in result
|
|
assert "created_at" in result
|
|
|
|
def test_create_collection_duplicate_name(self, clean_db):
|
|
"""Test creating collection with duplicate name."""
|
|
# Create first collection
|
|
clean_db.create_collection("duplicate_test", KBType.GENERIC)
|
|
|
|
# Try to create duplicate
|
|
with pytest.raises(ValueError, match="already exists"):
|
|
clean_db.create_collection("duplicate_test", KBType.CODE)
|
|
|
|
def test_create_collection_slug_conflict(self, clean_db):
|
|
"""Test creating collection with slug conflict."""
|
|
# Create first collection
|
|
clean_db.create_collection("Test Collection", KBType.GENERIC)
|
|
|
|
# Try to create collection with similar slug
|
|
with pytest.raises(ValueError, match="slug conflict"):
|
|
clean_db.create_collection("test-collection", KBType.CODE)
|
|
|
|
def test_list_collections(self, clean_db):
|
|
"""Test listing all collections."""
|
|
# Create a sample kb first
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
collections = clean_db.list_collections()
|
|
|
|
assert isinstance(collections, list)
|
|
assert len(collections) >= 1
|
|
assert any(col["name"] == kb_data["name"] for col in collections)
|
|
|
|
def test_get_collection_metadata(self, clean_db):
|
|
"""Test getting collection metadata."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
metadata = clean_db.get_collection_metadata(kb_data["name"])
|
|
|
|
assert metadata["kb_type"] == KBType.GENERIC.value
|
|
assert "description" in metadata
|
|
assert "created_at" in metadata
|
|
|
|
def test_get_collection_metadata_not_found(self, clean_db):
|
|
"""Test getting metadata for non-existent collection."""
|
|
with pytest.raises(ValueError, match="not found"):
|
|
clean_db.get_collection_metadata("nonexistent_collection")
|
|
|
|
def test_delete_collection(self, clean_db):
|
|
"""Test deleting a collection."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
# Verify collection exists
|
|
assert clean_db.collection_exists(kb_data["name"])
|
|
|
|
# Delete it
|
|
result = clean_db.delete_collection(kb_data["name"])
|
|
assert result is True
|
|
|
|
# Verify it's gone
|
|
assert not clean_db.collection_exists(kb_data["name"])
|
|
|
|
def test_delete_collection_not_found(self, clean_db):
|
|
"""Test deleting non-existent collection."""
|
|
result = clean_db.delete_collection("nonexistent_collection")
|
|
assert result is False
|
|
|
|
def test_collection_exists(self, clean_db):
|
|
"""Test checking if collection exists."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
assert clean_db.collection_exists(kb_data["name"]) is True
|
|
assert clean_db.collection_exists("nonexistent_collection") is False
|
|
|
|
def test_get_collection_count(self, clean_db):
|
|
"""Test getting document count for collection."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
# Add some documents
|
|
documents = ["Test document 1", "Test document 2"]
|
|
embeddings = [[0.1] * 768, [0.2] * 768]
|
|
metadatas = [{"filename": "doc1.txt"}, {"filename": "doc2.txt"}]
|
|
ids = ["doc1", "doc2"]
|
|
|
|
clean_db.add_documents(
|
|
kb_name=kb_data["name"],
|
|
documents=documents,
|
|
embeddings=embeddings,
|
|
metadatas=metadatas,
|
|
ids=ids
|
|
)
|
|
|
|
count = clean_db.get_collection_count(kb_data["name"])
|
|
assert count == 2
|
|
|
|
def test_get_collection_count_empty(self, clean_db):
|
|
"""Test getting document count for empty collection."""
|
|
# Create empty collection
|
|
clean_db.create_collection("empty_collection", KBType.GENERIC)
|
|
|
|
count = clean_db.get_collection_count("empty_collection")
|
|
assert count == 0
|
|
|
|
def test_get_collection_by_id(self, clean_db):
|
|
"""Test getting collection by ID."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
collection = clean_db.get_collection_by_id(kb_data["id"])
|
|
assert collection is not None
|
|
assert collection["name"] == kb_data["name"]
|
|
assert collection["id"] == kb_data["id"]
|
|
|
|
def test_get_collection_by_id_not_found(self, clean_db):
|
|
"""Test getting non-existent collection by ID."""
|
|
collection = clean_db.get_collection_by_id(99999)
|
|
assert collection is None
|
|
|
|
def test_list_collections_by_type(self, clean_db):
|
|
"""Test listing collections filtered by type."""
|
|
# Create collections of different types
|
|
clean_db.create_collection("generic_kb", KBType.GENERIC)
|
|
clean_db.create_collection("code_kb", KBType.CODE)
|
|
clean_db.create_collection("docs_kb", KBType.DOCUMENTATION)
|
|
|
|
# Get only code collections
|
|
code_collections = clean_db.list_collections_by_type(KBType.CODE)
|
|
assert len(code_collections) == 1
|
|
assert code_collections[0]["metadata"]["kb_type"] == "code"
|
|
|
|
def test_get_kb_by_slug(self, clean_db):
|
|
"""Test getting KB by slug."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
kb_name = clean_db.get_kb_by_slug(kb_data["slug"])
|
|
assert kb_name == kb_data["name"]
|
|
|
|
def test_get_kb_by_slug_not_found(self, clean_db):
|
|
"""Test getting non-existent KB by slug."""
|
|
kb_name = clean_db.get_kb_by_slug("nonexistent-slug")
|
|
assert kb_name is None
|
|
|
|
def test_slug_exists(self, clean_db):
|
|
"""Test checking if slug exists."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
assert clean_db.slug_exists(kb_data["slug"]) is True
|
|
assert clean_db.slug_exists("nonexistent-slug") is False
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestSQLiteManagerDocuments:
|
|
"""Test SQLite manager document operations."""
|
|
|
|
def test_add_documents(self, clean_db):
|
|
"""Test adding documents to collection."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
documents = ["Test document 1", "Test document 2"]
|
|
embeddings = [[0.1] * 768, [0.2] * 768]
|
|
metadatas = [{"filename": "doc1.txt"}, {"filename": "doc2.txt"}]
|
|
ids = ["doc1", "doc2"]
|
|
|
|
clean_db.add_documents(
|
|
kb_name=kb_data["name"],
|
|
documents=documents,
|
|
embeddings=embeddings,
|
|
metadatas=metadatas,
|
|
ids=ids
|
|
)
|
|
|
|
# Verify documents were added
|
|
count = clean_db.get_collection_count(kb_data["name"])
|
|
assert count == 2
|
|
|
|
def test_add_documents_nonexistent_collection(self, clean_db):
|
|
"""Test adding documents to non-existent collection."""
|
|
with pytest.raises(ValueError, match="not found"):
|
|
clean_db.add_documents(
|
|
kb_name="nonexistent_collection",
|
|
documents=["test"],
|
|
embeddings=[[0.1] * 768],
|
|
metadatas=[{}],
|
|
ids=["test"]
|
|
)
|
|
|
|
def test_get_documents_by_metadata(self, clean_db):
|
|
"""Test retrieving documents by metadata filter."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
# Add documents with metadata
|
|
for i in range(5):
|
|
np.random.seed(42 + i)
|
|
embedding = np.random.randn(768).tolist()
|
|
clean_db.add_documents(
|
|
kb_name=kb_data["name"],
|
|
documents=[f"This is test document {i}"],
|
|
embeddings=[embedding],
|
|
metadatas=[{"filename": f"test_{i}.txt", "chunk_index": i}],
|
|
ids=[f"node_doc_{i}"]
|
|
)
|
|
|
|
# Get all documents
|
|
all_docs = clean_db.get_documents_by_metadata(
|
|
kb_name=kb_data["name"],
|
|
where={}
|
|
)
|
|
assert len(all_docs) == 5
|
|
|
|
# Filter by filename
|
|
filtered_docs = clean_db.get_documents_by_metadata(
|
|
kb_name=kb_data["name"],
|
|
where={"filename": "test_0.txt"}
|
|
)
|
|
assert len(filtered_docs) == 1
|
|
assert filtered_docs[0]["metadata"]["filename"] == "test_0.txt"
|
|
|
|
def test_query_documents(self, clean_db):
|
|
"""Test querying documents by vector similarity."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
# Add documents
|
|
for i in range(5):
|
|
np.random.seed(42 + i)
|
|
embedding = np.random.randn(768).tolist()
|
|
clean_db.add_documents(
|
|
kb_name=kb_data["name"],
|
|
documents=[f"This is test document {i}"],
|
|
embeddings=[embedding],
|
|
metadatas=[{"filename": f"test_{i}.txt"}],
|
|
ids=[f"node_doc_{i}"]
|
|
)
|
|
|
|
query_embedding = [0.1] * 768
|
|
results = clean_db.query_documents(
|
|
kb_name=kb_data["name"],
|
|
query_embedding=query_embedding,
|
|
n_results=3
|
|
)
|
|
|
|
assert "ids" in results
|
|
assert "documents" in results
|
|
assert "metadatas" in results
|
|
assert "distances" in results
|
|
assert len(results["ids"][0]) <= 3
|
|
|
|
def test_query_documents_nonexistent_collection(self, clean_db):
|
|
"""Test querying non-existent collection."""
|
|
with pytest.raises(ValueError, match="not found"):
|
|
clean_db.query_documents(
|
|
kb_name="nonexistent_collection",
|
|
query_embedding=[0.1] * 768
|
|
)
|
|
|
|
def test_query_similar(self, clean_db):
|
|
"""Test querying similar documents by KB ID."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
# Add documents
|
|
for i in range(5):
|
|
np.random.seed(42 + i)
|
|
embedding = np.random.randn(768).tolist()
|
|
clean_db.add_documents(
|
|
kb_name=kb_data["name"],
|
|
documents=[f"This is test document {i}"],
|
|
embeddings=[embedding],
|
|
metadatas=[{"filename": f"test_{i}.txt"}],
|
|
ids=[f"node_doc_{i}"]
|
|
)
|
|
|
|
query_embedding = [0.1] * 768
|
|
results = clean_db.query_similar(
|
|
kb_id=kb_data["id"],
|
|
query_embedding=query_embedding,
|
|
top_k=3
|
|
)
|
|
|
|
assert isinstance(results, list)
|
|
assert len(results) <= 3
|
|
|
|
if results: # If we have results
|
|
for result in results:
|
|
assert "doc_id" in result
|
|
assert "text" in result
|
|
assert "metadata" in result
|
|
assert "similarity" in result
|
|
assert isinstance(result["similarity"], float)
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestSQLiteManagerProjects:
|
|
"""Test SQLite manager project operations."""
|
|
|
|
def test_create_project(self, clean_db):
|
|
"""Test creating a project."""
|
|
result = clean_db.create_project(
|
|
name="test_project",
|
|
path="/path/to/project",
|
|
description="Test project",
|
|
metadata={"type": "test"}
|
|
)
|
|
|
|
assert result["name"] == "test_project"
|
|
assert result["path"] == "/path/to/project"
|
|
assert result["description"] == "Test project"
|
|
assert result["metadata"]["type"] == "test"
|
|
assert "id" in result
|
|
assert "created_at" in result
|
|
|
|
def test_create_project_duplicate(self, clean_db):
|
|
"""Test creating duplicate project."""
|
|
clean_db.create_project("duplicate_project", "/path/1")
|
|
|
|
with pytest.raises(ValueError, match="already exists"):
|
|
clean_db.create_project("duplicate_project", "/path/2")
|
|
|
|
def test_get_project(self, clean_db):
|
|
"""Test getting project by ID."""
|
|
created = clean_db.create_project("test_project", "/path/to/project")
|
|
retrieved = clean_db.get_project(created["id"])
|
|
|
|
assert retrieved is not None
|
|
assert retrieved["name"] == "test_project"
|
|
assert retrieved["path"] == "/path/to/project"
|
|
|
|
def test_get_project_not_found(self, clean_db):
|
|
"""Test getting non-existent project."""
|
|
project = clean_db.get_project(99999)
|
|
assert project is None
|
|
|
|
def test_list_projects(self, clean_db):
|
|
"""Test listing all projects."""
|
|
# Create multiple projects
|
|
clean_db.create_project("project1", "/path/1")
|
|
clean_db.create_project("project2", "/path/2")
|
|
|
|
projects = clean_db.list_projects()
|
|
assert len(projects) >= 2
|
|
|
|
project_names = [p["name"] for p in projects]
|
|
assert "project1" in project_names
|
|
assert "project2" in project_names
|
|
|
|
def test_assign_kb_to_project(self, clean_db):
|
|
"""Test assigning KB to project."""
|
|
# Create project and KB
|
|
project = clean_db.create_project("test_project", "/path")
|
|
kb = clean_db.create_collection("test_kb", KBType.GENERIC)
|
|
|
|
# Assign KB to project
|
|
assignment = clean_db.assign_kb_to_project(
|
|
project_id=project["id"],
|
|
kb_id=kb["id"],
|
|
mcp_type="knowledge_docs"
|
|
)
|
|
|
|
assert assignment["project_id"] == project["id"]
|
|
assert assignment["kb_id"] == kb["id"]
|
|
assert assignment["mcp_type"] == "knowledge_docs"
|
|
|
|
def test_get_project_kbs(self, clean_db):
|
|
"""Test getting KB assignments for project."""
|
|
# Create project and KBs
|
|
project = clean_db.create_project("test_project", "/path")
|
|
kb1 = clean_db.create_collection("kb1", KBType.GENERIC)
|
|
kb2 = clean_db.create_collection("kb2", KBType.CODE)
|
|
|
|
# Assign KBs
|
|
clean_db.assign_kb_to_project(project["id"], kb1["id"], "knowledge_docs")
|
|
clean_db.assign_kb_to_project(project["id"], kb2["id"], "project_profile")
|
|
|
|
# Get assignments
|
|
kbs = clean_db.get_project_kbs(project["id"])
|
|
assert kbs["knowledge_docs"] == kb1["id"]
|
|
assert kbs["project_profile"] == kb2["id"]
|
|
|
|
def test_get_project_mcps_detailed(self, clean_db):
|
|
"""Test getting detailed MCP info for project."""
|
|
# Create project and KB
|
|
project = clean_db.create_project("test_project", "/path")
|
|
kb = clean_db.create_collection("test_kb", KBType.GENERIC)
|
|
|
|
# Assign KB
|
|
clean_db.assign_kb_to_project(project["id"], kb["id"], "knowledge_docs")
|
|
|
|
# Get detailed info
|
|
mcps = clean_db.get_project_mcps_detailed(project["id"])
|
|
assert "knowledge_docs" in mcps
|
|
assert mcps["knowledge_docs"]["kb_id"] == kb["id"]
|
|
assert mcps["knowledge_docs"]["kb_name"] == "test_kb"
|
|
|
|
def test_set_kb_folder(self, clean_db):
|
|
"""Test setting KB folder configuration."""
|
|
project = clean_db.create_project("test_project", "/path")
|
|
|
|
folder_config = clean_db.set_kb_folder(
|
|
project_id=project["id"],
|
|
mcp_type="knowledge_docs",
|
|
folder_path="/path/to/docs",
|
|
auto_sync=True
|
|
)
|
|
|
|
assert folder_config["project_id"] == project["id"]
|
|
assert folder_config["mcp_type"] == "knowledge_docs"
|
|
assert folder_config["folder_path"] == "/path/to/docs"
|
|
assert folder_config["auto_sync"] is True
|
|
|
|
def test_get_kb_folders(self, clean_db):
|
|
"""Test getting KB folder configurations."""
|
|
project = clean_db.create_project("test_project", "/path")
|
|
|
|
# Set multiple folder configs
|
|
clean_db.set_kb_folder(project["id"], "knowledge_docs", "/docs", True)
|
|
clean_db.set_kb_folder(project["id"], "project_profile", "/profile", False)
|
|
|
|
# Get all configs
|
|
folders = clean_db.get_kb_folders(project["id"])
|
|
assert folders["knowledge_docs"]["folder_path"] == "/docs"
|
|
assert folders["knowledge_docs"]["auto_sync"] is True
|
|
assert folders["project_profile"]["folder_path"] == "/profile"
|
|
assert folders["project_profile"]["auto_sync"] is False
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSQLiteManagerSingleton:
|
|
"""Test SQLite manager singleton pattern."""
|
|
|
|
def test_get_sqlite_manager_singleton(self):
|
|
"""Test that get_sqlite_manager returns singleton."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "test.db"
|
|
|
|
manager1 = get_sqlite_manager(str(db_path))
|
|
manager2 = get_sqlite_manager(str(db_path))
|
|
|
|
# Should be the same instance
|
|
assert manager1 is manager2
|
|
|
|
def test_get_sqlite_manager_default_path(self):
|
|
"""Test get_sqlite_manager with default path."""
|
|
# This should use the default path from Config
|
|
manager = get_sqlite_manager()
|
|
assert manager is not None
|
|
assert hasattr(manager, 'db_path')
|