"""Tests for prompt generation."""
import pytest
from code2prompt_rs import Code2Prompt
def test_generate_basic(test_dir):
"""Test basic prompt generation."""
prompt = Code2Prompt(path=test_dir)
result = prompt.generate()
# Basic checks
assert result.prompt is not None
assert isinstance(result.prompt, str)
assert result.token_count >= 0
assert str(result.directory) == test_dir
def test_generate_with_include_patterns(test_dir):
"""Test generation with include patterns."""
prompt = Code2Prompt(
path=test_dir,
include_patterns=["*.py"]
)
result = prompt.generate()
# Check that Python files are included
assert "foo.py" in result.prompt
assert "bar.py" in result.prompt
# Check that text files are excluded
assert "qux.txt" not in result.prompt
assert "corge.txt" not in result.prompt
def test_generate_with_exclude_patterns(test_dir):
"""Test generation with exclude patterns."""
prompt = Code2Prompt(
path=test_dir,
exclude_patterns=["**/uppercase/*"]
)
result = prompt.generate()
# Check that uppercase directory files are excluded
assert "FOO.py" not in result.prompt
assert "BAR.py" not in result.prompt
# Check that lowercase directory files are included
assert "foo.py" in result.prompt or "lowercase/foo.py" in result.prompt
def test_generate_with_line_numbers(test_dir):
"""Test generation with line numbers."""
prompt = Code2Prompt(
path=test_dir,
include_patterns=["lowercase/foo.py"],
line_numbers=True
)
result = prompt.generate()
# Check for line numbers in output (either format 1: or 1.|)
assert "1:" in result.prompt or "1 |" in result.prompt
def test_generate_with_relative_and_absolute_paths(test_dir):
"""Test generation with absolute paths."""
prompt_absolute = Code2Prompt(
path=test_dir,
include_patterns=["lowercase/foo.py"],
absolute_paths=True
)
result = prompt_absolute.generate()
# Should include absolute path format
assert test_dir in result.prompt
# Should include absolute path
assert "lowercase/foo.py" in result.prompt
prompt_relative = Code2Prompt(
path=test_dir,
include_patterns=["lowercase/foo.py"],
absolute_paths=False
)
result = prompt_relative.generate()
# Should not include absolute path format
assert test_dir not in result.prompt
# Should include absolute path
assert "lowercase/foo.py" in result.prompt
def test_generate_with_custom_template(test_dir):
"""Test generation with custom template."""
template = """# Code Overview
{% for file in files %}
## {{ file.path }}
```{{ file.language }}
{{ file.content }}" \
"{% endfor %}"""
prompt = Code2Prompt(
path=test_dir,
include_patterns=["lowercase/foo.py"]
)
result = prompt.generate(template=template)
# Check that custom template was used
assert "# Code Overview" in result.prompt
assert "## " in result.prompt
def test_token_count(test_dir):
"""Test token counting."""
prompt = Code2Prompt(path=test_dir)
# Get token count directly
token_count = prompt.token_count(encoding="cl100k")
assert isinstance(token_count, int)
assert token_count > 0
# Compare with generated result
result = prompt.generate(encoding="cl100k")
assert result.token_count == token_count
def test_multiple_encoding_options(test_dir):
"""Test with different encoding options."""
prompt = Code2Prompt(
path=test_dir,
include_patterns=["lowercase/foo.py"]
)
# Try different encodings
encodings = ["cl100k", "gpt2", "p50k_base"]
token_counts = {}
for encoding in encodings:
try:
count = prompt.token_count(encoding=encoding)
token_counts[encoding] = count
except Exception as e:
# Some encodings might not be available, that's OK
print(f"Encoding {encoding} failed: {e}")
# At least one encoding should work
assert len(token_counts) > 0
# Different encodings might give different counts
# (but for very small files they might be the same)
if len(token_counts) > 1:
unique_counts = set(token_counts.values())
print(f"Token counts: {token_counts}")
