File size: 2,355 Bytes
d61821a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | from __future__ import annotations
import unittest
from agent_harness.components import Candidate
from agent_harness.fusion import reciprocal_rank_fusion
from agent_harness.syntax_index import SyntaxRetriever, parse_go_file, parse_python_file
class SyntaxAndFusionTests(unittest.TestCase):
def test_tree_sitter_extracts_and_retrieves_method(self) -> None:
symbols = parse_go_file(
"cache/client.go",
"package cache\ntype Client struct{}\nfunc (c *Client) PresignHeadObject() error { return nil }\n",
)
names = {item.name for item in symbols}
self.assertEqual(names, {"Client", "PresignHeadObject"})
result = SyntaxRetriever(symbols).retrieve("HEAD presign cache object", 5)
self.assertEqual(result[0].symbol, "cache/client.go::PresignHeadObject")
def test_rrf_rewards_file_present_in_both_rankings(self) -> None:
def candidate(path: str, source: str) -> Candidate:
return Candidate(path, 1, 2, path, source, 1.0)
first = [candidate("a.go", "lexical"), candidate("shared.go", "lexical")]
second = [candidate("b.go", "dense"), candidate("shared.go", "dense")]
fused = reciprocal_rank_fusion([first, second], 3)
self.assertEqual(fused[0].path, "shared.go")
def test_tree_sitter_extracts_and_retrieves_python_method(self) -> None:
symbols = parse_python_file(
"gitlab/client.py",
"class Client:\n def retry_request(self, url: str) -> bool:\n return bool(url)\n",
)
names = {item.name for item in symbols}
self.assertEqual(names, {"Client", "retry_request"})
result = SyntaxRetriever(symbols).retrieve("retry client request", 5)
self.assertEqual(result[0].symbol, "gitlab/client.py::retry_request")
self.assertIn("retry_request", symbols[1].identifiers)
def test_identifier_extraction_persists_after_first_declaration(self) -> None:
symbols = parse_go_file(
"demo.go",
"package demo\nfunc First() { Second() }\nfunc Second() { First() }\n",
)
self.assertEqual([item.name for item in symbols], ["First", "Second"])
self.assertIn("Second", symbols[0].identifiers)
self.assertIn("First", symbols[1].identifiers)
if __name__ == "__main__":
unittest.main()
|