File size: 2,852 Bytes
27ad088 |
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 |
import pytest
import tempfile
from unittest.mock import Mock, patch, AsyncMock
from pathlib import Path
from pydantic_ai import Agent
from know_lang_bot.code_parser.summarizer import CodeSummarizer
from know_lang_bot.code_parser.parser import CodeChunk, ChunkType
from know_lang_bot.config import AppConfig
@pytest.fixture
def config():
"""Create a test configuration"""
with tempfile.TemporaryDirectory() as temp_dir:
yield AppConfig(
llm={"model_name": "test-model", "model_provider": "test"},
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('know_lang_bot.code_parser.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('know_lang_bot.code_parser.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 |