1 files changed, 37 insertions, 0 deletions
diff --git a/continuedev/src/continuedev/libs/steps/pytest.py b/continuedev/src/continuedev/libs/steps/pytest.py
new file mode 100644
index 00000000..e53eb465
--- /dev/null
+++ b/continuedev/src/continuedev/libs/steps/pytest.py
@@ -0,0 +1,37 @@
+from textwrap import dedent
+from ...models.filesystem_edit import AddDirectory, AddFile
+from ..core import Step, ContinueSDK
+import os
+
+
+class WritePytestsStep(Step):
+    for_filepath: str
+
+    async def run(self, sdk: ContinueSDK):
+        filename, dirname = os.path.split(self.for_filepath)
+
+        path_dir = os.path.join(dirname, "tests")
+        if not os.path.exists(path_dir):
+            await sdk.apply_filesystem_edit(AddDirectory(path=path_dir))
+
+        path = os.path.join(path_dir, f"test_{filename}")
+        if os.path.exists(path):
+            return
+
+        for_file_contents = await sdk.ide.readFile(self.for_filepath)
+
+        prompt = dedent(f"""\
+            This is the file you will write unit tests for:
+
+            ```python
+            {for_file_contents}
+            ```
+
+            Here are additional instructions:
+
+            "{self.instructions}"
+
+            Here is a complete set of pytest unit tests:
+        """)
+        tests = sdk.llm.complete(prompt)
+        await sdk.apply_filesystem_edit(AddFile(filepath=path, content=tests))