Created a new patchflow that can generate unit tests (#998)

* initial changes

* Patched /home/priyash7/Desktop/open-source/patchwork/patchwork/patchflows/GenerateUnitTests/README.md

* final changes

* intermediate changes

* made changes to modifyCode step

* added README and fixed code

* partition changes

* added test to CI

* added default extension to yaml

---------

Co-authored-by: patched.codes[bot] <298395+patched.codes[bot]@users.noreply.github.com>

Author: plon-Susk7

.github/workflows/test.yml

@@ -8,6 +8,7 @@ on:
       - dependencyupgrade-*
       - generatereadme-*
       - generatedocstring-*
+      - generateunittests-*
       - demo*
 
 jobs:
@@ -201,6 +202,14 @@ jobs:
           --base_path=tests/cicd/generate_docstring \
           --disable_telemetry
 
+      - name: Generate UnitTests
+        run: |
+          poetry run patchwork GenerateUnitTests --log debug \
+          --patched_api_key=${{ secrets.PATCHED_API_KEY }} \
+          --github_api_key=${{ secrets.SCM_GITHUB_KEY }} \
+          --folder_path=tests/cicd/generate_docstring \
+          --disable_telemetry
+
       - name: Generate README
         run: |
           # Specify the parent folder you want to check

patchwork/patchflows/GenerateREADME/README.md

@@ -92,4 +92,4 @@ You can update the default [prompt template](./generate_readme_prompt.json). Not
 Here are some example PRs generated with the GenerateREADME patchflow:
 
 - https://github.com/patched-codes/patchwork/pull/53
-- https://github.com/patched-codes/patchwork/pull/52
+- https://github.com/patched-codes/patchwork/pull/52

patchwork/patchflows/GenerateUnitTests/GenerateUnitTests.py

@@ -0,0 +1,61 @@
+import json
+from pathlib import Path
+import yaml
+
+from patchwork.common.utils.step_typing import validate_steps_with_inputs
+from patchwork.step import Step
+from patchwork.steps import (
+    LLM, 
+    CallCode2Prompt,
+    ModifyCode,
+    PR
+)
+
+_DEFAULT_INPUT_FILE = Path(__file__).parent / "defaults.yml"
+_DEFAULT_PROMPT_JSON = Path(__file__).parent / "default_prompt.json"
+
+class GenerateUnitTests(Step):
+    def __init__(self, inputs):
+        super().__init__(inputs)
+
+        final_inputs = yaml.safe_load(_DEFAULT_INPUT_FILE.read_text())
+        if final_inputs is None:
+            final_inputs = {}
+        
+        final_inputs.update(inputs)
+
+        final_inputs["prompt_id"] = "GenerateUnitTests"
+        if "folder_path" not in final_inputs.keys():
+            final_inputs["folder_path"] = Path.cwd()
+        else:
+            final_inputs["folder_path"] = Path(final_inputs["folder_path"])
+
+        if "prompt_template_file" not in final_inputs:
+            final_inputs["prompt_template_file"] = _DEFAULT_PROMPT_JSON
+
+        final_inputs["pr_title"] = f"PatchWork Unit Tests generated"
+        final_inputs["branch_prefix"] = f"{self.__class__.__name__.lower()}-"
+
+        validate_steps_with_inputs(
+            set(final_inputs.keys()).union({"prompt_values","files_to_patch"}), LLM, CallCode2Prompt,ModifyCode,PR
+        )
+        self.inputs = final_inputs
+
+    def run(self):
+        outputs = CallCode2Prompt(self.inputs).run()
+        new_file_name = f"test_file.{self.inputs['test_file_extension']}"
+        new_file_path = Path(outputs['uri']).with_name(new_file_name)
+        Path(outputs['uri']).rename(new_file_path)
+        outputs['uri'] = str(new_file_path)
+        self.inputs["response_partitions"] = {"patch": ["```", "\n", "```"]}
+        self.inputs["files_to_patch"] = self.inputs["prompt_values"] = [outputs]
+        outputs = LLM(self.inputs).run()
+        self.inputs.update(outputs)
+        outputs = ModifyCode(self.inputs).run()
+        self.inputs.update(outputs)
+        number = len(self.inputs["modified_code_files"])
+        self.inputs["pr_header"] = f"This pull request from patchwork adds tests."
+        outputs = PR(self.inputs).run()
+        self.inputs.update(outputs)
+
+        return self.inputs

patchwork/patchflows/GenerateUnitTests/README.md

@@ -0,0 +1,12 @@
+## Code Documentation
+
+### Inputs
+- The code reads default inputs from a YAML file (`defaults.yml`) and a JSON file (`default_prompt.json`).
+- The code takes user inputs and updates the default inputs accordingly.
+- The code expects inputs like `folder_path`, `prompt_template_file`, `test_file_extension`, etc.
+
+### Outputs
+- The code generates unit tests based on the provided inputs.
+- It creates a new test file with the specified extension.
+- The code updates various parameters in the inputs dictionary during its execution.
+- The final output includes the modified inputs after running the unit test generation process.

patchwork/patchflows/GenerateUnitTests/__init__.py

@@ -0,0 +1,15 @@
+[
+    {
+        "id": "GenerateUnitTests", 
+        "prompts": [
+            {
+                "role": "system",
+                "content": "You are a senior software tester who is highly skilled at writing unit tests across various programming languages. Users will specify classes or functions to test, and you will generate appropriate unit tests using the most relevant framework for the detected language. Output the code directly with no additional text, formatting, or triple quotes. The response should appear as if directly pasted into an editor."
+            },
+            {
+                "role": "user",
+                "content": "Code: {{fullContent}}"
+            }
+        ]
+    }
+]

patchwork/patchflows/GenerateUnitTests/default_prompt.json

@@ -0,0 +1,25 @@
+# CallLLM Inputs
+# openai_api_key: required-for-chatgpt
+# google_api_key: required-for-gemini
+# model: gpt-4o
+# client_base_url: https://api.openai.com/v1
+# Example HF model
+# client_base_url: https://api-inference.huggingface.co/models/codellama/CodeLlama-70b-Instruct-hf/v1
+# model: codellama/CodeLlama-70b-Instruct-hf
+# model_temperature: 0.2
+# model_top_p: 0.95
+# model_max_tokens: 2000
+
+# folder_path : path/to/folder/with/class
+
+# Default value
+test_file_extension : py
+
+# CommitChanges Inputs
+disable_branch: false
+
+# CreatePR Inputs
+disable_pr: false
+force_pr_creation: true
+# github_api_key: required-for-github-scm
+# gitlab_api_key: required-for-gitlab-scm

patchwork/patchflows/GenerateUnitTests/defaults.yml

@@ -4,5 +4,6 @@
 from .GenerateREADME.GenerateREADME import GenerateREADME
 from .PRReview.PRReview import PRReview
 from .ResolveIssue.ResolveIssue import ResolveIssue
+from .GenerateUnitTests.GenerateUnitTests import GenerateUnitTests
 
-__all__ = ["AutoFix", "DependencyUpgrade", "GenerateREADME", "PRReview", "ResolveIssue", "GenerateDocstring"]
+__all__ = ["AutoFix", "DependencyUpgrade", "GenerateREADME", "PRReview", "ResolveIssue", "GenerateDocstring", "GenerateUnitTests"]

patchwork/patchflows/__init__.py

@@ -85,9 +85,10 @@ def run(self) -> dict:
             start_line = code_snippet.get("startLine")
             end_line = code_snippet.get("endLine")
             new_code = extracted_response.get("patch")
+
             if new_code is None:
                 continue
-
+            
             replace_code_in_file(uri, start_line, end_line, new_code)
             modified_code_file = dict(path=uri, start_line=start_line, end_line=end_line, **extracted_response)
             modified_code_files.append(modified_code_file)