From b5d58ab2c7f7d1f5b4d375e4f1696bc038cb070e Mon Sep 17 00:00:00 2001 From: Harshad Hegde Date: Mon, 12 Feb 2024 12:03:25 -0600 Subject: [PATCH 1/9] Added `write-tests` feature --- src/codergpt/__init__.py | 1 + src/codergpt/cli.py | 23 +++++++++ src/codergpt/constants.py | 1 + src/codergpt/main.py | 10 ++-- src/codergpt/test_writer/test_writer.py | 69 +++++++++++++++++++++++++ src/codergpt/test_writer/tester.py | 1 - tests/input/math.py | 33 ++++++++++++ tests/test_tester.py | 68 ++++++++++++++++++++++++ tox.ini | 1 + 9 files changed, 202 insertions(+), 5 deletions(-) create mode 100644 src/codergpt/test_writer/test_writer.py delete mode 100644 src/codergpt/test_writer/tester.py create mode 100644 tests/input/math.py create mode 100644 tests/test_tester.py diff --git a/src/codergpt/__init__.py b/src/codergpt/__init__.py index dff9f87..433b590 100644 --- a/src/codergpt/__init__.py +++ b/src/codergpt/__init__.py @@ -6,6 +6,7 @@ from codergpt.explainer import CodeExplainer from codergpt.optimizer import CodeOptimizer +# from codergpt.tester import CodeTester from .main import CoderGPT try: diff --git a/src/codergpt/cli.py b/src/codergpt/cli.py index 1f0df4f..7231abc 100644 --- a/src/codergpt/cli.py +++ b/src/codergpt/cli.py @@ -146,5 +146,28 @@ def optimize_code(path: Union[str, Path], function: str, classname: str, overwri raise ValueError("The path provided is not a file.") +@main.command("write-tests") +@path_argument +@function_option +@class_option +def write_test_code(path: Union[str, Path], function: str, classname: str): + """ + Write tests for the code file. + + :param path: The path to the code file. + :param function: The name of the function to test. Default is None. + :param classname: The name of the class to test. Default is None. + """ + # Ensure path is a string or Path object for consistency + if isinstance(path, str): + path = Path(path) + + # Check if path is a file + if path.is_file(): + coder.test_writer(path=path, function=function, classname=classname) + else: + raise ValueError("The path provided is not a file.") + + if __name__ == "__main__": main() diff --git a/src/codergpt/constants.py b/src/codergpt/constants.py index ad46ddb..7e63676 100644 --- a/src/codergpt/constants.py +++ b/src/codergpt/constants.py @@ -2,6 +2,7 @@ from pathlib import Path +TEST_DIR = Path(__file__).resolve().parents[2] / "tests" SRC = Path(__file__).resolve().parents[1] PACKAGE_DIR = SRC / "codergpt" EXTENSION_MAP_FILE = PACKAGE_DIR / "extensions.yaml" diff --git a/src/codergpt/main.py b/src/codergpt/main.py index aced52b..9749897 100644 --- a/src/codergpt/main.py +++ b/src/codergpt/main.py @@ -10,15 +10,16 @@ from tabulate import tabulate from codergpt.commenter.commenter import CodeCommenter -from codergpt.constants import EXTENSION_MAP_FILE, GPT_3_5_TURBO, INSPECTION_HEADERS +from codergpt.constants import EXTENSION_MAP_FILE, GPT_4_TURBO, INSPECTION_HEADERS from codergpt.explainer.explainer import CodeExplainer from codergpt.optimizer.optimizer import CodeOptimizer +from codergpt.test_writer.test_writer import CodeTester class CoderGPT: """CoderGPT class.""" - def __init__(self, model: str = GPT_3_5_TURBO): + def __init__(self, model: str = GPT_4_TURBO): """Initialize the CoderGPT class.""" self.llm = ChatOpenAI(openai_api_key=os.environ.get("OPENAI_API_KEY"), temperature=0.7, model=model) self.prompt = ChatPromptTemplate.from_messages( @@ -126,13 +127,14 @@ def optimizer(self, path: Union[str, Path], function: str = None, classname=None # code, language = self.get_code(filename=path, function_name=function, class_name=classname) code_optimizer.optimize(filename=path, function=function, classname=classname, overwrite=overwrite) - def tester(self, path: Union[str, Path]): + def test_writer(self, path: Union[str, Path], function: str = None, classname: str = None): """ Test the code file. :param path: The path to the code file. """ - pass + code_tester = CodeTester(self.chain) + code_tester.write_tests(filename=path, function=function, classname=classname) if __name__ == "__main__": diff --git a/src/codergpt/test_writer/test_writer.py b/src/codergpt/test_writer/test_writer.py new file mode 100644 index 0000000..ca226c9 --- /dev/null +++ b/src/codergpt/test_writer/test_writer.py @@ -0,0 +1,69 @@ +"""Test writing module.""" + +import os +from pathlib import Path +from typing import Any, Dict, Optional, Union + +from langchain_core.runnables.base import RunnableSerializable + +from codergpt.constants import TEST_DIR + + +class CodeTester: + """Code tester class writes testing code from a given file.""" + + def __init__(self, chain: RunnableSerializable[Dict, Any]): + """ + Initialize the CodeTester class with a runnable chain. + + :param chain: A RunnableSerializable object capable of executing tasks. + """ + self.chain = chain + + def write_tests( + self, + filename: Union[str, Path], + function: Optional[str] = None, + classname: Optional[str] = None, + outfile: Optional[str] = None, + ): + """ + Write tests for the code by invoking the runnable chain. + + :param path: The path to the code file to be explained. + :param function: The name of the function to explain. Default is None. + :param classname: The name of the class to explain. Default is None. + """ + with open(filename, "r") as source_file: + source_code = source_file.read() + if function: + response = self.chain.invoke( + { + "input": f"Write tests for the function '{function}' in \n\n```\n{source_code}\n```" + "Return just the code block. Also explain the tests in a systematic way as a comment." + } + ) + elif classname: + response = self.chain.invoke( + { + "input": f"Write tests for the class '{classname}' in \n\n```\n{source_code}\n```" + "Also explain the tests in a systematic way." + } + ) + else: + # Write tests for full code + response = self.chain.invoke( + { + "input": f"Write tests for the following code: \n\n```\n{source_code}\n```" + "Also explain the tests in a systematic way." + } + ) + test_code = response.content + base_filename = os.path.basename(filename) + if outfile: + new_filepath = outfile + else: + new_filepath = f"{TEST_DIR}/test_{base_filename}" + # Write the test to the new file + with open(new_filepath, "w") as updated_file: + updated_file.write(test_code) diff --git a/src/codergpt/test_writer/tester.py b/src/codergpt/test_writer/tester.py deleted file mode 100644 index 5cb2099..0000000 --- a/src/codergpt/test_writer/tester.py +++ /dev/null @@ -1 +0,0 @@ -"""Test writing module.""" diff --git a/tests/input/math.py b/tests/input/math.py new file mode 100644 index 0000000..81d23ac --- /dev/null +++ b/tests/input/math.py @@ -0,0 +1,33 @@ +"""Test python code.""" + +def calculate_sum(numbers): + """ + Calculate the sum of a list of numbers. + + :param numbers: A list of numbers. + :type numbers: list[int] + :return: The sum of the numbers. + :rtype: int + """ + result = 0 + for number in numbers: + result += number + return result + + +class MathOperations: + def multiply(self, a, b): + """ + Multiply two numbers. + + :param a: The first number. + :type a: int + :param b: The second number. + :type b: int + :return: The product of the two numbers. + :rtype: int + """ + answer = 0 + for i in range(b): + answer += a + return answer diff --git a/tests/test_tester.py b/tests/test_tester.py new file mode 100644 index 0000000..1aa63b9 --- /dev/null +++ b/tests/test_tester.py @@ -0,0 +1,68 @@ +"""Tests for the CodeTester class.""" + +import os +import unittest +from pathlib import Path + +from codergpt.test_writer.test_writer import CodeTester +from langchain_core.prompts import ChatPromptTemplate +from langchain_openai import ChatOpenAI + +from .test_constants import TEST_INPUT_DIR, TEST_OUTPUT_DIR + + +class CodeTesterTests(unittest.TestCase): + """Tests for the CodeTester class.""" + + def setUp(self): + """Create a sample runnable chain for testing.""" + self.llm = ChatOpenAI(openai_api_key=os.environ.get("OPENAI_API_KEY")) + self.prompt = ChatPromptTemplate.from_messages( + [("system", "You are world class software developer."), ("user", "{input}")] + ) + self.chain = self.prompt | self.llm + self.code_tester = CodeTester(chain=self.chain) + self.filename = TEST_INPUT_DIR / "math.py" + self.output_filename = TEST_OUTPUT_DIR / "test_math.py" + + def tearDown(self): + """Clean up the created test files.""" + test_files = Path(TEST_OUTPUT_DIR).glob("test_*") + for file in test_files: + file.unlink() + + @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skipping tests in GitHub Actions") + def test_write_tests_with_function(self): + """Test writing tests for a function.""" + # Arrange + function = "calculate_sum" + + # Act + outfile = TEST_OUTPUT_DIR / "test_math_function.py" + self.code_tester.write_tests(self.filename, function=function, outfile=outfile) + + self.assertTrue(outfile.exists()) + + @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skipping tests in GitHub Actions") + def test_write_tests_with_class(self): + """Test writing tests for a class.""" + # Arrange + classname = "MathOperations" + + # Act + outfile = TEST_OUTPUT_DIR / "test_math_class.py" + self.code_tester.write_tests(self.filename, classname=classname, outfile=outfile) + + self.assertTrue(outfile.exists()) + + @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skipping tests in GitHub Actions") + def test_write_tests_without_function_or_class(self): + """Test writing tests for a file.""" + # Act + outfile = self.output_filename + self.code_tester.write_tests(self.filename, outfile=outfile) + self.assertTrue(outfile.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tox.ini b/tox.ini index affc57b..5a3b0d5 100644 --- a/tox.ini +++ b/tox.ini @@ -21,6 +21,7 @@ envlist = [testenv] allowlist_externals = poetry +passenv = OPENAI_API_KEY commands = poetry run pytest {posargs} description = Run unit tests with pytest. This is a special environment that does not get a name, and From b7f6eb14272a1bc870f00a671a2c684307620a57 Mon Sep 17 00:00:00 2001 From: Harshad Hegde Date: Mon, 12 Feb 2024 12:04:14 -0600 Subject: [PATCH 2/9] added docstring --- tests/input/math.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/input/math.py b/tests/input/math.py index 81d23ac..5429cea 100644 --- a/tests/input/math.py +++ b/tests/input/math.py @@ -16,6 +16,8 @@ def calculate_sum(numbers): class MathOperations: + """Class to perform mathematical operations.""" + def multiply(self, a, b): """ Multiply two numbers. From 3a167629b2a2f05a8e55320230b17e510c381288 Mon Sep 17 00:00:00 2001 From: Harshad Hegde Date: Mon, 12 Feb 2024 12:18:33 -0600 Subject: [PATCH 3/9] Added documentation --- README.md | 44 +++++++++++++++++++++++ docs/description.rst | 46 +++++++++++++++++++++++++ docs/test.rst | 30 ++++++++++++++++ src/codergpt/optimizer/optimizer.py | 33 +++++++++--------- src/codergpt/test_writer/test_writer.py | 29 ++++++++-------- 5 files changed, 152 insertions(+), 30 deletions(-) create mode 100644 docs/test.rst diff --git a/README.md b/README.md index 69ead52..d56e02e 100644 --- a/README.md +++ b/README.md @@ -192,6 +192,50 @@ code [OPTIONS] COMMAND [ARGS]... """ ``` +5. `write-tests`: Writes tests for the specified code file. The user can specify a function and/or a class within the file to target with the tests. + +```shell +code write-tests [--function ] [--class ] +``` + +#### Example +- Let's consider a python file `example.py`: +```python +# example.py + +def add(a, b): + return a + b + +class Calculator: + def subtract(self, a, b): + return a - b +``` +```shell +$ code write-tests example.py --function add --class Calculator +``` +results in test files being generated that contain test cases for the `add` function and the `Calculator` class. The actual content of the test files will depend on the implementation of the `coder.test_writer` method but would typically look something like this: + +```python +import unittest +from example import add, Calculator + +class TestAddFunction(unittest.TestCase): + + def test_addition(self): + self.assertEqual(add(3, 4), 7) + +class TestCalculator(unittest.TestCase): + + def setUp(self): + self.calc = Calculator() + + def test_subtract(self): + self.assertEqual(self.calc.subtract(10, 5), 5) +``` + +In this example, running the command generates unit tests for both the `add` function and the `Calculator` class in the `example.py` file. The tests check if the `add` function correctly adds two numbers and if the `Calculator`'s `subtract` method correctly subtracts one number from another. + + ## Development The CLI is built using Python and the `click` library. Below is an example of how to define a new command: diff --git a/docs/description.rst b/docs/description.rst index 1e83cb1..97232e5 100644 --- a/docs/description.rst +++ b/docs/description.rst @@ -202,6 +202,52 @@ Commands By using these optimizations, we improve the efficiency and readability of the code. """ +5s. **write-tests**: Generates test cases for specified functions and/or classes within a Python code file. + + .. code-block:: shell + + code write-tests [--function ] [--class ] + + **Example** + + - Let's consider a Python file `example.py`: + + .. code-block:: python + + # example.py + + def add(a, b): + return a + b + + class Calculator: + def subtract(self, a, b): + return a - b + + .. code-block:: shell + + $ code write-tests example.py --function add --class Calculator + + results in the creation of test files that contain test cases for both the `add` function and the `Calculator` class. The content of the generated test files might look like this: + + .. code-block:: python + + import unittest + from example import add, Calculator + + class TestAddFunction(unittest.TestCase): + + def test_addition(self): + self.assertEqual(add(3, 4), 7) + + class TestCalculator(unittest.TestCase): + + def setUp(self): + self.calc = Calculator() + + def test_subtract(self): + self.assertEqual(self.calc.subtract(10, 5), 5) + + In this example, executing the command generates unit tests for the `add` function and the `Calculator` class defined in `example.py`. The tests verify whether the `add` function correctly computes the sum of two numbers and if the `Calculator`'s `subtract` method accurately performs subtraction. Development ----------- diff --git a/docs/test.rst b/docs/test.rst new file mode 100644 index 0000000..0361e96 --- /dev/null +++ b/docs/test.rst @@ -0,0 +1,30 @@ +.. py:module:: codergpt + +Test writing module +=================== + +.. py:class:: CodeTester(chain) + + The CodeTester class is responsible for generating testing code from a given source file. It utilizes a llm chain to produce tests for specific functions or classes within the source file. + + .. py:method:: __init__(chain) + + Initializes the CodeTester instance with a provided llm chain. + + :param chain: A RunnableSerializable object capable of executing tasks. + :type chain: RunnableSerializable[Dict, Any] + + .. py:method:: write_tests(filename, function=None, classname=None, outfile=None) + + Generates test cases for the specified code by invoking the llm chain. If a function or class name is provided, it will generate tests specifically for that function or class. Otherwise, it will attempt to create tests for the entire code. + + :param filename: The path to the code file for which tests are to be written. + :type filename: Union[str, Path] + :param function: The name of the function for which tests should be generated. Defaults to None, indicating that no specific function is targeted. + :type function: Optional[str] + :param classname: The name of the class for which tests should be generated. Defaults to None, indicating that no specific class is targeted. + :type classname: Optional[str] + :param outfile: The path where the generated test file should be saved. If not provided, a default path within the TEST_DIR will be used. + :type outfile: Optional[str] + + The method reads the source code from the provided filename and uses the llm chain to generate appropriate test cases. The resulting test code is then written to either the specified outfile or a new file within the TEST_DIR directory. diff --git a/src/codergpt/optimizer/optimizer.py b/src/codergpt/optimizer/optimizer.py index 379366b..c5a02d3 100644 --- a/src/codergpt/optimizer/optimizer.py +++ b/src/codergpt/optimizer/optimizer.py @@ -33,22 +33,23 @@ def optimize( """ with open(filename, "r") as source_file: source_code = source_file.read() - if function: - response = self.chain.invoke( - { - "input": f"Optimize, comment and add sphinx docstrings" - f" to the function '{function}' in \n\n```\n{source_code}\n```" - "Also explain the optimization in a systematic way as a comment." - } - ) - elif classname: - response = self.chain.invoke( - { - "input": f"Optimize, comment and add sphinx docstrings" - f" to the class '{classname}' in \n\n```\n{source_code}\n```" - "Also explain the optimization in a systematic way." - } - ) + if function or classname: + if function: + response = self.chain.invoke( + { + "input": f"Optimize, comment and add sphinx docstrings" + f" to the function '{function}' in \n\n```\n{source_code}\n```" + "Also explain the optimization in a systematic way as a comment." + } + ) + if classname: + response = self.chain.invoke( + { + "input": f"Optimize, comment and add sphinx docstrings" + f" to the class '{classname}' in \n\n```\n{source_code}\n```" + "Also explain the optimization in a systematic way." + } + ) else: # Optimize full code response = self.chain.invoke( diff --git a/src/codergpt/test_writer/test_writer.py b/src/codergpt/test_writer/test_writer.py index ca226c9..1e01289 100644 --- a/src/codergpt/test_writer/test_writer.py +++ b/src/codergpt/test_writer/test_writer.py @@ -36,20 +36,21 @@ def write_tests( """ with open(filename, "r") as source_file: source_code = source_file.read() - if function: - response = self.chain.invoke( - { - "input": f"Write tests for the function '{function}' in \n\n```\n{source_code}\n```" - "Return just the code block. Also explain the tests in a systematic way as a comment." - } - ) - elif classname: - response = self.chain.invoke( - { - "input": f"Write tests for the class '{classname}' in \n\n```\n{source_code}\n```" - "Also explain the tests in a systematic way." - } - ) + if function or classname: + if function: + response = self.chain.invoke( + { + "input": f"Write tests for the function '{function}' in \n\n```\n{source_code}\n```" + "Return just the code block. Also explain the tests in a systematic way as a comment." + } + ) + if classname: + response = self.chain.invoke( + { + "input": f"Write tests for the class '{classname}' in \n\n```\n{source_code}\n```" + "Also explain the tests in a systematic way." + } + ) else: # Write tests for full code response = self.chain.invoke( From deea045accb35b784c07db2e4c4f0ed97d259500 Mon Sep 17 00:00:00 2001 From: Harshad Hegde Date: Mon, 12 Feb 2024 12:23:37 -0600 Subject: [PATCH 4/9] Added skip test ig gh workflow --- tests/test_tester.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_tester.py b/tests/test_tester.py index 1aa63b9..f1489f9 100644 --- a/tests/test_tester.py +++ b/tests/test_tester.py @@ -10,7 +10,7 @@ from .test_constants import TEST_INPUT_DIR, TEST_OUTPUT_DIR - +@unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skipping tests in GitHub Actions") class CodeTesterTests(unittest.TestCase): """Tests for the CodeTester class.""" From 8cd28ba332186f9eb5796019ac26c06ab0874d3e Mon Sep 17 00:00:00 2001 From: Harshad Hegde Date: Mon, 12 Feb 2024 12:23:49 -0600 Subject: [PATCH 5/9] formatted --- tests/test_tester.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_tester.py b/tests/test_tester.py index f1489f9..15bb5af 100644 --- a/tests/test_tester.py +++ b/tests/test_tester.py @@ -10,6 +10,7 @@ from .test_constants import TEST_INPUT_DIR, TEST_OUTPUT_DIR + @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skipping tests in GitHub Actions") class CodeTesterTests(unittest.TestCase): """Tests for the CodeTester class.""" From 5de8382c31f926d921fb62e47179e27457244bcb Mon Sep 17 00:00:00 2001 From: Harshad Hegde Date: Mon, 12 Feb 2024 12:25:50 -0600 Subject: [PATCH 6/9] Added skip test ig gh workflow --- tests/test_tester.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_tester.py b/tests/test_tester.py index 15bb5af..6b4b12b 100644 --- a/tests/test_tester.py +++ b/tests/test_tester.py @@ -11,10 +11,10 @@ from .test_constants import TEST_INPUT_DIR, TEST_OUTPUT_DIR -@unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skipping tests in GitHub Actions") class CodeTesterTests(unittest.TestCase): """Tests for the CodeTester class.""" + @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skipping tests in GitHub Actions") def setUp(self): """Create a sample runnable chain for testing.""" self.llm = ChatOpenAI(openai_api_key=os.environ.get("OPENAI_API_KEY")) @@ -26,6 +26,7 @@ def setUp(self): self.filename = TEST_INPUT_DIR / "math.py" self.output_filename = TEST_OUTPUT_DIR / "test_math.py" + @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skipping tests in GitHub Actions") def tearDown(self): """Clean up the created test files.""" test_files = Path(TEST_OUTPUT_DIR).glob("test_*") From f37399f7d64c5970ee76ed10e293cda08291e17f Mon Sep 17 00:00:00 2001 From: Harshad Hegde Date: Mon, 12 Feb 2024 12:31:21 -0600 Subject: [PATCH 7/9] trying to skip test file in a workflow --- .github/workflows/qc.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/qc.yml b/.github/workflows/qc.yml index a75049c..06ebac8 100644 --- a/.github/workflows/qc.yml +++ b/.github/workflows/qc.yml @@ -33,4 +33,4 @@ jobs: run: poetry run tox -e lint - name: Test with pytest and generate coverage file - run: poetry run tox -e py + run: poetry run tox -e py -- --ignore=tests/test_tester.py From 9bdd47b1e0820d9705619d06c0a599531c33e14e Mon Sep 17 00:00:00 2001 From: Harshad Hegde Date: Mon, 12 Feb 2024 12:33:34 -0600 Subject: [PATCH 8/9] >=3.9 --- .github/workflows/qc.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/qc.yml b/.github/workflows/qc.yml index 06ebac8..2dcf262 100644 --- a/.github/workflows/qc.yml +++ b/.github/workflows/qc.yml @@ -11,7 +11,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [ "3.8", "3.11" ] + python-version: [ "3.9", "3.11" ] steps: - uses: actions/checkout@v3.0.2 From a313b5b056ac174bff6c71d54e4eb086270c739c Mon Sep 17 00:00:00 2001 From: Harshad Hegde Date: Mon, 12 Feb 2024 12:40:05 -0600 Subject: [PATCH 9/9] pytest instead of tox --- .github/workflows/qc.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/qc.yml b/.github/workflows/qc.yml index 2dcf262..5be1464 100644 --- a/.github/workflows/qc.yml +++ b/.github/workflows/qc.yml @@ -33,4 +33,5 @@ jobs: run: poetry run tox -e lint - name: Test with pytest and generate coverage file - run: poetry run tox -e py -- --ignore=tests/test_tester.py + run: poetry run pytest --ignore=tests/test_tester.py + # Changed from tox to pytest to make --ignore work.