File size: 2,235 Bytes
e8c208d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
from typing import Dict, Any
from flow_modules.Tachi67.CodeFileEditFlowModule import CodeFileEditAtomicFlow

class TestCodeFileEditAtomicFlow(CodeFileEditAtomicFlow):
    def _generate_import_statement(self, code_lib_location):
        module_dir = os.path.dirname(code_lib_location)
        module_name = os.path.basename(code_lib_location).rstrip('.py')
        import_code = (
            f"import sys\n"
            f"sys.path.insert(0, '{module_dir}')\n"
            f"import {module_name}\n"
        )
        return import_code

    def _generate_content(self, code_lib_location, code_str) -> str:
        import_code_lib_str = self._generate_import_statement(code_lib_location)
        content = (
                "# Don't touch this import statement \n"
                + import_code_lib_str + "\n"
                "# Here is the code just generated \n" +
                code_str + "\n"
                "# Below, please provide code to test it.\n"
                "# The simplest form could be just calling it with appropriate parameters. \n"
                "# You could also assert the output, anyway, the test results will be informed to JARVIS. \n"
                "# If you do not write anything, JARVIS just checks if the syntax is alright. \n"
                "###########\n"
                "\n # Test Code:\n" +
                "\n############\n"
        )
        return content

    def _generate_temp_file_location(self, code_lib_location):
        directory = os.path.dirname(code_lib_location)
        ret = os.path.join(directory, 'temp_tests.py')
        return ret

    def _check_input(self, input_data: Dict[str, Any]):
        assert "code" in input_data, "code is not passed to TestCodeFileEditAtomicFlow"
        assert "code_library" in input_data, "code_library is not passed to TestCodeFileEditAtomicFlow"

    def _generate_input_to_writer(self, input_data: Dict[str, Any]):
        code_str = input_data['code']
        code_lib_location = input_data["code_library"]
        content_to_write = self._generate_content(code_lib_location, code_str)
        file_location_to_write = self._generate_temp_file_location(code_lib_location)
        return content_to_write, file_location_to_write