diff --git a/tests/test_formatter.py b/tests/test_formatter.py new file mode 100644 index 0000000..6d8bdab --- /dev/null +++ b/tests/test_formatter.py @@ -0,0 +1,71 @@ +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