76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
from app.modules.rag.contracts.enums import RagLayer
|
|
from app.modules.rag.indexing.code.pipeline import CodeIndexingPipeline
|
|
|
|
|
|
def test_code_pipeline_builds_source_symbols_edges_and_entrypoints() -> None:
|
|
pipeline = CodeIndexingPipeline()
|
|
content = """
|
|
from fastapi import APIRouter
|
|
|
|
router = APIRouter()
|
|
|
|
class UserService:
|
|
def get_user(self, user_id):
|
|
return user_id
|
|
|
|
@router.get("/users/{user_id}")
|
|
async def get_user(user_id: str):
|
|
service = UserService()
|
|
return service.get_user(user_id)
|
|
"""
|
|
docs = pipeline.index_file(
|
|
repo_id="acme/proj",
|
|
commit_sha="abc123",
|
|
path="app/api/users.py",
|
|
content=content,
|
|
)
|
|
|
|
layers = {doc.layer for doc in docs}
|
|
assert RagLayer.CODE_SOURCE_CHUNKS in layers
|
|
assert RagLayer.CODE_SYMBOL_CATALOG in layers
|
|
assert RagLayer.CODE_DEPENDENCY_GRAPH in layers
|
|
assert RagLayer.CODE_ENTRYPOINTS in layers
|
|
|
|
symbol_doc = next(doc for doc in docs if doc.layer == RagLayer.CODE_SYMBOL_CATALOG and doc.metadata["kind"] == "function")
|
|
assert "get_user" in symbol_doc.metadata["qname"]
|
|
|
|
edge_doc = next(doc for doc in docs if doc.layer == RagLayer.CODE_DEPENDENCY_GRAPH)
|
|
assert edge_doc.metadata["edge_type"] in {"calls", "imports", "inherits"}
|
|
|
|
entry_doc = next(doc for doc in docs if doc.layer == RagLayer.CODE_ENTRYPOINTS)
|
|
assert entry_doc.metadata["framework"] == "fastapi"
|
|
|
|
|
|
def test_code_pipeline_indexes_import_alias_as_symbol() -> None:
|
|
pipeline = CodeIndexingPipeline()
|
|
content = "from .v2 import ConfigManagerV2 as ConfigManager\n"
|
|
|
|
docs = pipeline.index_file(
|
|
repo_id="acme/proj",
|
|
commit_sha="abc123",
|
|
path="src/config_manager/__init__.py",
|
|
content=content,
|
|
)
|
|
|
|
alias_doc = next(doc for doc in docs if doc.layer == RagLayer.CODE_SYMBOL_CATALOG and doc.metadata["qname"] == "ConfigManager")
|
|
assert alias_doc.metadata["kind"] == "const"
|
|
assert alias_doc.metadata["lang_payload"]["import_alias"] is True
|
|
|
|
|
|
def test_code_pipeline_marks_test_documents() -> None:
|
|
pipeline = CodeIndexingPipeline()
|
|
content = """
|
|
def test_user_service():
|
|
assert True
|
|
"""
|
|
|
|
docs = pipeline.index_file(
|
|
repo_id="acme/proj",
|
|
commit_sha="abc123",
|
|
path="tests/test_users.py",
|
|
content=content,
|
|
)
|
|
|
|
assert docs
|
|
assert all(doc.metadata["is_test"] is True for doc in docs)
|