test: wrote tests for context tool

Author: ErikBjare

gptme/tools/context.py

@@ -37,50 +37,52 @@ def _gitignore():
     return ignored
 
 
-def _ctags():
+def _ctags() -> str:
     """Generate ctags for current project."""
 
     ignored = _gitignore()
     ignored_str = " ".join([f"--exclude='{i}'" for i in ignored])
 
     shell = get_shell()
     cmd = f"ctags -R --output-format=json {ignored_str} --fields=+l+n --languages=python --python-kinds=-iv -f -"
-    print(cmd)
     ret, ctags, _ = shell.run_command(cmd)
     assert ret == 0
 
-    print("ctags:")
+    output = ["ctags:"]
     tags = []
     for line in ctags.splitlines():
         try:
             tags.append(json.loads(line))
         except json.JSONDecodeError:
-            print("  failed to parse: ", line)
+            output += ["  failed to parse: {line}"]
             break
 
     files = {tag["path"] for tag in tags}
     for file in sorted(files):
         filetags = [tag for tag in tags if tag["path"] == file]
-        print(f"{file}")
+        output += [str(file)]
         level = 0
         for tag in sorted(filetags, key=lambda x: x["line"]):
             if tag["kind"] == "class":
-                print(level * "  " + f"  class {tag['name']}:{tag['line']}")
+                output += [level * "  " + f"  class {tag['name']}:{tag['line']}"]
                 level += 1
             elif tag["kind"] == "function":
                 level = 0
-                print(level * "  " + f"  def {tag['name']}:{tag['line']}")
+                output += [level * "  " + f"  def {tag['name']}:{tag['line']}"]
             elif tag["kind"] == "variable":
                 level = 0
-                print(level * "  " + f"  {tag['name']}:{tag['line']}")
+                output += [level * "  " + f"  {tag['name']}:{tag['line']}"]
             elif tag["kind"] == "unknown":
                 # import a as b
                 pass
             else:
-                print(level * "  " + f"  {tag['kind']} {tag['name']}:{tag['line']}")
+                output += [
+                    level * "  " + f"  {tag['kind']} {tag['name']}:{tag['line']}"
+                ]
 
-    return ctags
+    return "\n".join(output)
 
 
 if __name__ == "__main__":
-    assert _ctags()
+    output = _ctags()
+    print(output)

gptme/tools/shell.py

@@ -23,6 +23,9 @@ def __init__(self) -> None:
         self.stderr_fd = self.process.stderr.fileno()  # type: ignore
         self.delimiter = "END_OF_COMMAND_OUTPUT"
 
+        # close on exit
+        atexit.register(self.close)
+
     def run_command(self, command: str) -> tuple[int | None, str, str]:
         assert self.process.stdin
 
@@ -80,12 +83,15 @@ def get_shell() -> ShellSession:
     if _shell is None:
         # init shell
         _shell = ShellSession()
-
-        # close on exit
-        atexit.register(_shell.close)
     return _shell
 
 
+# used in testing
+def set_shell(shell: ShellSession) -> None:
+    global _shell
+    _shell = shell
+
+
 def execute_shell(cmd: str, ask=True) -> Generator[Message, None, None]:
     """Executes a shell command and returns the output."""
     shell = get_shell()

tests/test_tools_context.py

@@ -0,0 +1,21 @@
+import pytest
+from gptme.tools.context import _ctags, _gen_context_msg
+from gptme.tools.shell import ShellSession, set_shell
+
+
+@pytest.fixture
+def shell():
+    shell = ShellSession()
+    set_shell(shell)
+    return shell
+
+
+def test_gen_context_msg(shell):
+    msg = _gen_context_msg()
+    assert "gptme" in msg.content, f"Expected 'gptme' in output: {msg.content}"
+    assert "$ pwd" in msg.content, f"Expected 'pwd' in output: {msg.content}"
+
+
+def test_ctags(shell):
+    output = _ctags()
+    assert "function" in output, f"Expected 'def' in output: {output}"