blob: 63edabc6d646c474dc7a8ec2dcb5d4cd2ee997ac (
plain)
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
|
import os
from textwrap import dedent
from typing import Union
from ....core.main import Step
from ....core.sdk import ContinueSDK
from ....models.filesystem_edit import AddDirectory, AddFile
class WritePytestsRecipe(Step):
for_filepath: Union[str, None] = None
user_input: str = "Write unit tests for this file."
async def describe(self, models):
return f"Writing unit tests for {self.for_filepath}"
async def run(self, sdk: ContinueSDK):
if self.for_filepath is None:
self.for_filepath = (await sdk.ide.getVisibleFiles())[0]
filename = os.path.basename(self.for_filepath)
dirname = os.path.dirname(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 None
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.user_input}"
Here is a complete set of pytest unit tests:"""
)
tests = await sdk.models.summarize.complete(prompt)
await sdk.apply_filesystem_edit(AddFile(filepath=path, content=tests))
return None
|
