Merge pull request #34 from biobootloader/improve-structure-add-tests

Improve project structure and add tests

README.md

@@ -23,11 +23,11 @@ _warning!_ By default wolverine uses GPT-4 and may make many repeated calls to t
 
 To run with gpt-4 (the default, tested option):
 
-    python wolverine.py examples/buggy_script.py "subtract" 20 3
+    python -m wolverine examples/buggy_script.py "subtract" 20 3
 
 You can also run with other models, but be warned they may not adhere to the edit format as well:
 
-    python wolverine.py --model=gpt-3.5-turbo buggy_script.py "subtract" 20 3
+    python -m wolverine --model=gpt-3.5-turbo examples/buggy_script.py "subtract" 20 3
 
 If you want to use GPT-3.5 by default instead of GPT-4 uncomment the default model line in `.env`:
 
@@ -35,7 +35,7 @@ If you want to use GPT-3.5 by default instead of GPT-4 uncomment the default mod
 
 You can also use flag `--confirm=True` which will ask you `yes or no` before making changes to the file. If flag is not used then it will apply the changes to the file
 
-    python wolverine.py buggy_script.py "subtract" 20 3 --confirm=True
+    python -m wolverine examples/buggy_script.py "subtract" 20 3 --confirm=True
 
 ## Future Plans
 

requirements-dev.txt

@@ -0,0 +1,2 @@
+-r requirements.txt
+pytest==7.3.1

tests/test_wolverine.py

@@ -0,0 +1,56 @@
+import os
+import json
+import pytest
+import tempfile
+from wolverine import apply_changes, json_validated_response
+
+
+@pytest.fixture(scope='function')
+def temp_file():
+    # Create a temporary file
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
+        f.write("first line\nsecond line\nthird line")
+    file_path = f.name
+    yield file_path
+    # Clean up the temporary file
+    os.remove(file_path)
+
+
+def test_apply_changes_replace(temp_file):
+    # Make a "replace" change to the second line
+    changes = [
+        {"operation": "Replace", "line": 2, "content": "new second line"}
+    ]
+    apply_changes(temp_file, changes)
+
+    # Check that the file was updated correctly
+    with open(temp_file) as f:
+        content = f.read()
+        assert content == "first line\nnew second line\nthird line"
+
+
+def test_apply_changes_delete(temp_file):
+    # Make a "delete" change to the third line
+    changes = [
+        {"operation": "Delete", "line": 3, "content": ""},
+    ]
+    apply_changes(temp_file, changes)
+
+    # Check that the file was updated correctly
+    with open(temp_file) as f:
+        content = f.read()
+        assert content == "first line\nsecond line\n"
+
+
+def test_apply_changes_insert(temp_file):
+    # Make an "insert" change after the second line
+    changes = [
+        {"operation": "InsertAfter", "line": 2, "content": "inserted line"},
+    ]
+    apply_changes(temp_file, changes)
+
+    # Check that the file was updated correctly
+    with open(temp_file) as f:
+        content = f.read()
+        assert content == 'first line\nsecond line\ninserted line\nthird line'
+

wolverine/__init__.py

@@ -0,0 +1 @@
+from .wolverine import apply_changes, json_validated_response

wolverine/__main__.py

@@ -0,0 +1,6 @@
+import fire
+
+from .wolverine import main
+
+if __name__ == "__main__":
+    fire.Fire(main)

wolverine/wolverine.py

@@ -1,5 +1,4 @@
 import difflib
-import fire
 import json
 import os
 import shutil
@@ -20,6 +19,7 @@ DEFAULT_MODEL = os.environ.get("DEFAULT_MODEL", "gpt-4")
 with open("prompt.txt") as f:
     SYSTEM_PROMPT = f.read()
 
+
 def run_script(script_name, script_args):
     script_args = [str(arg) for arg in script_args]
     """
@@ -224,8 +224,3 @@ def main(script_name, *script_args, revert=False, model=DEFAULT_MODEL, confirm=F
 
             apply_changes(script_name, json_response, confirm=confirm)
             cprint("Changes applied. Rerunning...", "blue")
-
-
-if __name__ == "__main__":
-    fire.Fire(main)
-