72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
import pytest
|
|
from pathlib import Path
|
|
from codechunk.core.formatter import OutputFormatter
|
|
from codechunk.core.chunking import ParsedChunk, ChunkMetadata
|
|
|
|
|
|
class TestOutputFormatter:
|
|
"""Tests for OutputFormatter."""
|
|
|
|
def test_format_markdown(self, sample_python_code, tmp_path):
|
|
"""Test markdown output format."""
|
|
formatter = OutputFormatter(format_type="markdown")
|
|
|
|
from codechunk.core.parser import CodeParser
|
|
parser = CodeParser()
|
|
file_path = tmp_path / "test.py"
|
|
file_path.write_text(sample_python_code)
|
|
|
|
chunks = parser.parse_file(file_path)
|
|
output = formatter.format(chunks)
|
|
|
|
assert "# Code Context" in output
|
|
assert len(output) > 0
|
|
|
|
def test_format_ollama(self, sample_python_code, tmp_path):
|
|
"""Test Ollama output format."""
|
|
formatter = OutputFormatter(format_type="ollama")
|
|
|
|
from codechunk.core.parser import CodeParser
|
|
parser = CodeParser()
|
|
file_path = tmp_path / "test.py"
|
|
file_path.write_text(sample_python_code)
|
|
|
|
chunks = parser.parse_file(file_path)
|
|
output = formatter.format(chunks)
|
|
|
|
assert "### System" in output
|
|
assert "### User" in output
|
|
|
|
def test_format_lmstudio(self, sample_python_code, tmp_path):
|
|
"""Test LM Studio output format."""
|
|
formatter = OutputFormatter(format_type="lmstudio")
|
|
|
|
from codechunk.core.parser import CodeParser
|
|
parser = CodeParser()
|
|
file_path = tmp_path / "test.py"
|
|
file_path.write_text(sample_python_code)
|
|
|
|
chunks = parser.parse_file(file_path)
|
|
output = formatter.format(chunks)
|
|
|
|
assert output.strip().startswith("[")
|
|
assert output.strip().endswith("]")
|
|
|
|
def test_estimate_tokens(self):
|
|
"""Test token estimation."""
|
|
formatter = OutputFormatter()
|
|
|
|
text = "This is a test string with about twenty words in it"
|
|
tokens = formatter.estimate_tokens(text)
|
|
|
|
assert tokens > 0
|
|
assert tokens < len(text)
|
|
|
|
def test_empty_chunks(self):
|
|
"""Test formatting with no chunks."""
|
|
formatter = OutputFormatter(format_type="markdown")
|
|
|
|
output = formatter.format([])
|
|
|
|
assert len(output) > 0
|