add

patchwork/common/utils.py

@@ -16,7 +16,7 @@
 def defered_temp_file(
     mode="w+b", buffering=-1, encoding=None, newline=None, suffix=None, prefix=None, dir=None, *, errors=None
 ):
-    tempfile_fp = tempfile.TemporaryFile(
+    tempfile_fp = tempfile.NamedTemporaryFile(
         mode=mode,
         buffering=buffering,
         encoding=encoding,

patchwork/steps/PreparePR/PreparePR.py

@@ -15,7 +15,7 @@ def __init__(self, inputs: dict):
             raise ValueError(f'Missing required data: "{self.required_keys}"')
 
         if len(inputs["modified_code_files"]) < 1:
-            logger.warn("No modified files to prepare a PR for.")
+            logger.warning("No modified files to prepare a PR for.")
         self.modified_code_files = inputs["modified_code_files"]
 
         self.header = f"This pull request from patched fixes {len(self.modified_code_files)} issues."

tests/steps/test_ExtractCode.py

@@ -0,0 +1,86 @@
+import json
+import os
+from pathlib import Path
+
+import pytest
+
+from patchwork.steps.ExtractCode.ExtractCode import ExtractCode, Severity
+
+_DEFAULT_SARIF_FILE_NAME = "sarif_file.sarif"
+
+
+@pytest.fixture
+def extract_code_instance(tmp_path):
+    original_dir = Path.cwd()
+
+    os.chdir(tmp_path)
+
+    test_file = tmp_path / "test.py"
+    test_file.write_text("print('Hello, world!')")
+
+    sarif_file = tmp_path / _DEFAULT_SARIF_FILE_NAME
+    sarif_file.write_text(
+        json.dumps(
+            {
+                "runs": [
+                    {
+                        "results": [
+                            {
+                                "message": {"text": "Error message"},
+                                "ruleId": "1",
+                                "locations": [
+                                    {
+                                        "physicalLocation": {
+                                            "artifactLocation": {"uri": str(test_file)},
+                                            "region": {"startLine": 1, "endLine": 1},
+                                        }
+                                    }
+                                ],
+                            }
+                        ],
+                        "tool": {"driver": {"rules": [{"id": "1", "defaultConfiguration": {"level": "high"}}]}},
+                    }
+                ],
+            }
+        )
+    )
+    inputs = {
+        "sarif_file_path": sarif_file,
+        "context_size": 1000,
+        "vulnerability_limit": 10,
+        "severity": "HIGH",
+    }
+    yield ExtractCode(inputs)
+    os.chdir(original_dir)
+
+
+def test_extract_code_init(extract_code_instance):
+    assert extract_code_instance.sarif_file_path.name == _DEFAULT_SARIF_FILE_NAME
+    assert extract_code_instance.context_length == 1000
+    assert extract_code_instance.vulnerability_limit == 10
+    assert extract_code_instance.severity_threshold == Severity.HIGH
+
+
+def test_extract_code_run(extract_code_instance, tmp_path):
+    # Run the extract code step
+    result = extract_code_instance.run()
+
+    # Check that the extracted code contexts are correct
+    assert len(extract_code_instance.extracted_code_contexts) == 1
+    assert extract_code_instance.extracted_code_contexts[0]["uri"] == "test.py"
+    assert extract_code_instance.extracted_code_contexts[0]["startLine"] == 0
+    assert extract_code_instance.extracted_code_contexts[0]["endLine"] == 1
+    assert extract_code_instance.extracted_code_contexts[0]["affectedCode"] == "print('Hello, world!')"
+    assert extract_code_instance.extracted_code_contexts[0]["messageText"] == "Error message"
+
+    # Check that the output file is created
+    assert result["code_file"].exists()
+    assert result["code_file"].is_file()
+    with open(result["code_file"], "r") as f:
+        output_data = json.load(f)
+        assert len(output_data) == 1
+        assert output_data[0]["uri"] == "test.py"
+        assert output_data[0]["startLine"] == 0
+        assert output_data[0]["endLine"] == 1
+        assert output_data[0]["affectedCode"] == "print('Hello, world!')"
+        assert output_data[0]["messageText"] == "Error message"

tests/steps/test_ExtractModelResponse.py

@@ -0,0 +1,38 @@
+import pytest
+
+from patchwork.steps.ExtractModelResponse.ExtractModelResponse import (
+    ExtractModelResponse,
+)
+
+
+@pytest.fixture
+def sample_inputs():
+    return {
+        "openai_responses": ["partition1response1partition2", "response2partition3"],
+        "response_partitions": {"key1": ["partition1", "partition2"], "key2": ["partition3"]},
+    }
+
+
+def test_init_required_keys(sample_inputs):
+    step = ExtractModelResponse(sample_inputs)
+    assert step.openai_responses == sample_inputs["openai_responses"]
+    assert step.partitions == sample_inputs["response_partitions"]
+
+
+def test_init_missing_required_keys():
+    with pytest.raises(ValueError):
+        ExtractModelResponse({})
+
+
+def test_run_no_partitions(sample_inputs):
+    step = ExtractModelResponse({**sample_inputs, "response_partitions": {}})
+    output = step.run()
+    assert output == {"extracted_responses": []}
+
+
+def test_run_with_partitions(sample_inputs):
+    step = ExtractModelResponse(sample_inputs)
+    output = step.run()
+    assert len(output["extracted_responses"]) == 2
+    assert output["extracted_responses"][0]["key1"] == "response1"
+    assert output["extracted_responses"][1]["key2"] == "response2"

tests/steps/test_GenerateEmbeddings.py

@@ -0,0 +1,44 @@
+import pytest
+
+from patchwork.steps.GenerateEmbeddings.GenerateEmbeddings import (
+    GenerateEmbeddings,
+    filter_by_extension,
+)
+
+
+def test_filter_by_extension(tmp_path):
+    assert filter_by_extension("example.txt", [".txt"])
+    assert not filter_by_extension("example.txt", [".pdf"])
+
+
+def test_generate_embeddings_init():
+    inputs = {"embedding_name": "test", "documents": [{"document": "test document"}]}
+    step = GenerateEmbeddings(inputs)
+    assert step.collection is not None
+    assert step.documents == inputs["documents"]
+
+
+def test_generate_embeddings_run():
+    inputs = {"embedding_name": "test", "documents": [{"document": "test document"}]}
+    step = GenerateEmbeddings(inputs)
+    result = step.run()
+    assert result == {}
+
+
+def test_generate_embeddings_init_required_keys_missing():
+    inputs = {"documents": [{"document": "test document"}]}
+    with pytest.raises(ValueError):
+        GenerateEmbeddings(inputs)
+
+
+def test_generate_embeddings_init_embedding_name_missing():
+    inputs = {"embedding_name": "test"}
+    with pytest.raises(ValueError):
+        GenerateEmbeddings(inputs)
+
+
+def test_generate_embeddings_run():
+    inputs = {"embedding_name": "test", "documents": [{"document": "test document"}]}
+    step = GenerateEmbeddings(inputs)
+    result = step.run()
+    assert result == {}