106 lines
3.2 KiB
Python
106 lines
3.2 KiB
Python
import pytest
|
|
import tree_sitter
|
|
from file_cache import ASTParser
|
|
|
|
def test_ast_parser_initialization():
|
|
"""Verify that ASTParser can be initialized with a language string."""
|
|
parser = ASTParser("python")
|
|
assert parser.language_name == "python"
|
|
|
|
def test_ast_parser_parse():
|
|
"""Verify that the parse method returns a tree_sitter.Tree."""
|
|
parser = ASTParser("python")
|
|
code = """def example_func():
|
|
return 42"""
|
|
tree = parser.parse(code)
|
|
assert isinstance(tree, tree_sitter.Tree)
|
|
# Basic check that it parsed something
|
|
assert tree.root_node.type == "module"
|
|
|
|
def test_ast_parser_get_skeleton_python():
|
|
"""Verify that get_skeleton replaces function bodies with '...' while preserving docstrings."""
|
|
parser = ASTParser("python")
|
|
code = '''
|
|
def complex_function(a, b):
|
|
"""
|
|
This is a docstring.
|
|
It should be preserved.
|
|
"""
|
|
result = a + b
|
|
if result > 0:
|
|
return result
|
|
return 0
|
|
|
|
class MyClass:
|
|
def method_without_docstring(self):
|
|
print("doing something")
|
|
return None
|
|
'''
|
|
skeleton = parser.get_skeleton(code)
|
|
|
|
# Check that signatures are preserved
|
|
assert "def complex_function(a, b):" in skeleton
|
|
assert "class MyClass:" in skeleton
|
|
assert "def method_without_docstring(self):" in skeleton
|
|
|
|
# Check that docstring is preserved
|
|
assert '"""' in skeleton
|
|
assert "This is a docstring." in skeleton
|
|
assert "It should be preserved." in skeleton
|
|
|
|
# Check that bodies are replaced with '...'
|
|
assert "..." in skeleton
|
|
assert "result = a + b" not in skeleton
|
|
assert "return result" not in skeleton
|
|
assert 'print("doing something")' not in skeleton
|
|
|
|
def test_ast_parser_invalid_language():
|
|
"""Verify handling of unsupported or invalid languages."""
|
|
# This might raise an error or return a default, depending on implementation
|
|
# For now, we expect it to either fail gracefully or raise an exception we can catch
|
|
with pytest.raises(Exception):
|
|
ASTParser("not-a-language")
|
|
|
|
def test_ast_parser_get_curated_view():
|
|
"""Verify that get_curated_view preserves function bodies with @core_logic or # [HOT]."""
|
|
parser = ASTParser("python")
|
|
code = '''
|
|
@core_logic
|
|
def core_func():
|
|
"""Core logic doc."""
|
|
print("this should be preserved")
|
|
return True
|
|
|
|
def hot_func():
|
|
# [HOT]
|
|
print("this should also be preserved")
|
|
return 42
|
|
|
|
def normal_func():
|
|
"""Normal doc."""
|
|
print("this should be stripped")
|
|
return None
|
|
|
|
class MyClass:
|
|
@core_logic
|
|
def core_method(self, x):
|
|
print("method preserved", x)
|
|
'''
|
|
curated = parser.get_curated_view(code)
|
|
|
|
# Check that core_func is preserved
|
|
assert 'print("this should be preserved")' in curated
|
|
assert 'return True' in curated
|
|
|
|
# Check that hot_func is preserved
|
|
assert '# [HOT]' in curated
|
|
assert 'print("this should also be preserved")' in curated
|
|
|
|
# Check that normal_func is stripped but docstring is preserved
|
|
assert '"""Normal doc."""' in curated
|
|
assert 'print("this should be stripped")' not in curated
|
|
assert '...' in curated
|
|
|
|
# Check that core_method is preserved
|
|
assert 'print("method preserved", x)' in curated
|