File size: 4,774 Bytes
27ad088 60532a1 27ad088 3536fc0 27ad088 60532a1 27ad088 60532a1 27ad088 369993f 60532a1 369993f 3536fc0 369993f 3536fc0 369993f 3536fc0 369993f 410a99b 3536fc0 369993f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 |
import pytest
import tempfile
from unittest.mock import Mock, patch, AsyncMock
from pathlib import Path
from knowlang.summarizer.summarizer import CodeSummarizer
from knowlang.core.types import CodeChunk, ChunkType
from knowlang.configs.config import AppConfig
@pytest.fixture
def config():
"""Create a test configuration"""
with tempfile.TemporaryDirectory() as temp_dir:
yield AppConfig(
llm={"model_name": "testing", "model_provider": "testing"},
db={"persist_directory": Path(temp_dir), "collection_name": "test_collection"}
)
@pytest.fixture
def sample_chunks():
"""Create sample code chunks for testing"""
return [
CodeChunk(
type=ChunkType.FUNCTION,
content="def hello(): return 'world'",
start_line=1,
end_line=2,
file_path="test.py",
name="hello",
docstring="Says hello"
),
CodeChunk(
type=ChunkType.CLASS,
content="class TestClass:\n def __init__(self):\n pass",
start_line=4,
end_line=6,
file_path="test.py",
name="TestClass",
docstring="A test class"
)
]
@pytest.fixture
def mock_summary():
"""Create a sample summary result"""
return "This is a test function"
@pytest.fixture
def mock_run_result(mock_summary):
"""Create a mock run result"""
mock_result = Mock()
mock_result.data = mock_summary
return mock_result
@pytest.mark.asyncio
@patch('knowlang.summarizer.summarizer.Agent')
async def test_summarize_chunk(mock_agent_class, config: AppConfig, sample_chunks: list[CodeChunk], mock_run_result: Mock):
"""Test summarizing a single chunk"""
# Setup the mock agent instance
mock_agent = mock_agent_class.return_value
mock_agent.run = AsyncMock(return_value=mock_run_result)
summarizer = CodeSummarizer(config)
result = await summarizer.summarize_chunk(sample_chunks[0])
# Verify result
assert isinstance(result, str)
assert result == mock_run_result.data
# Verify agent was called with correct prompt
call_args = mock_agent.run.call_args[0][0]
assert "def hello()" in call_args
assert "Says hello" in call_args
@patch('knowlang.summarizer.summarizer.Agent')
def test_chromadb_initialization(mock_agent_class, config: AppConfig):
"""Test ChromaDB initialization"""
mock_agent = mock_agent_class.return_value
summarizer = CodeSummarizer(config)
assert summarizer.collection is not None
# Verify we can create a new collection
summarizer.db_client.delete_collection(config.db.collection_name)
new_summarizer = CodeSummarizer(config)
assert new_summarizer.collection is not None
@pytest.mark.asyncio
@patch('knowlang.summarizer.summarizer.generate_embedding')
@patch('knowlang.summarizer.summarizer.Agent')
async def test_process_and_store_chunk_with_embedding(
mock_agent_class,
mock_embedding_generator,
config: AppConfig,
sample_chunks: list[CodeChunk],
mock_run_result: Mock
):
"""Test processing and storing a chunk with embedding"""
# Setup the mock agent instance
mock_agent = mock_agent_class.return_value
mock_agent.run = AsyncMock(return_value=mock_run_result)
# Setup mock embedding response
mock_embedding = [0.1, 0.2, 0.3] # Sample embedding vector
mock_embedding_generator.return_value = mock_embedding
summarizer = CodeSummarizer(config)
# Mock the collection's add method
summarizer.collection.add = Mock()
# Process the chunk
await summarizer.process_and_store_chunk(sample_chunks[0])
# Verify ollama.embed was called with correct parameters
mock_embedding_generator.assert_called_once_with(
mock_run_result.data,
config.embedding,
)
# Verify collection.add was called with correct parameters
add_call = summarizer.collection.add.call_args
assert add_call is not None
kwargs = add_call[1]
assert len(kwargs['embeddings']) == 3
assert kwargs['embeddings'] == mock_embedding
assert kwargs['documents'][0] == mock_run_result.data
assert kwargs['ids'][0] == f"{sample_chunks[0].file_path}:{sample_chunks[0].start_line}-{sample_chunks[0].end_line}"
# Verify metadata
metadata = kwargs['metadatas'][0]
assert metadata['file_path'] == sample_chunks[0].file_path
assert metadata['start_line'] == sample_chunks[0].start_line
assert metadata['end_line'] == sample_chunks[0].end_line
assert metadata['type'] == sample_chunks[0].type.value
assert metadata['name'] == sample_chunks[0].name
assert metadata['docstring'] == sample_chunks[0].docstring |