Adding strict-coverage option and removing prompt only. (#78)

Codium-ai · Jun 3, 2024 · 3be1457 · 3be1457
1 parent 5685c9f
commit 3be1457
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 42 deletions.
diff --git a/cover_agent/CoverAgent.py b/cover_agent/CoverAgent.py
@@ -1,6 +1,7 @@
 import datetime
 import os
 import shutil
+import sys
 import wandb
 
 from cover_agent.CustomLogger import CustomLogger
@@ -47,59 +48,57 @@ def _duplicate_test_file(self):
             self.args.test_file_output_path = self.args.test_file_path
 
     def run(self):
-
         if 'WANDB_API_KEY' in os.environ:
             wandb.login(key=os.environ['WANDB_API_KEY'])
             time_and_date = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
             run_name = f"{self.args.model}_" + time_and_date
             wandb.init(project="cover-agent", name=run_name)
 
-        if not self.args.prompt_only:
-            iteration_count = 0
-            test_results_list = []
-
-            self.test_gen.initial_test_suite_analysis()
+        iteration_count = 0
+        test_results_list = []
 
-            while (
-                self.test_gen.current_coverage < (self.test_gen.desired_coverage / 100)
-                and iteration_count < self.args.max_iterations
-            ):
-                self.logger.info(
-                    f"Current Coverage: {round(self.test_gen.current_coverage * 100, 2)}%"
-                )
-                self.logger.info(f"Desired Coverage: {self.test_gen.desired_coverage}%")
+        self.test_gen.initial_test_suite_analysis()
 
-                generated_tests_dict = self.test_gen.generate_tests(max_tokens=4096)
+        while (
+            self.test_gen.current_coverage < (self.test_gen.desired_coverage / 100)
+            and iteration_count < self.args.max_iterations
+        ):
+            self.logger.info(
+                f"Current Coverage: {round(self.test_gen.current_coverage * 100, 2)}%"
+            )
+            self.logger.info(f"Desired Coverage: {self.test_gen.desired_coverage}%")
 
-                for generated_test in generated_tests_dict.get("new_tests", []):
-                    test_result = self.test_gen.validate_test(
-                        generated_test, generated_tests_dict
-                    )
-                    test_results_list.append(test_result)
+            generated_tests_dict = self.test_gen.generate_tests(max_tokens=4096)
 
-                iteration_count += 1
+            for generated_test in generated_tests_dict.get("new_tests", []):
+                test_result = self.test_gen.validate_test(
+                    generated_test, generated_tests_dict
+                )
+                test_results_list.append(test_result)
 
-                if self.test_gen.current_coverage < (
-                    self.test_gen.desired_coverage / 100
-                ):
-                    self.test_gen.run_coverage()
+            iteration_count += 1
 
-            if self.test_gen.current_coverage >= (self.test_gen.desired_coverage / 100):
-                self.logger.info(
-                    f"Reached above target coverage of {self.test_gen.desired_coverage}% (Current Coverage: {round(self.test_gen.current_coverage * 100, 2)}%) in {iteration_count} iterations."
-                )
-            elif iteration_count == self.args.max_iterations:
-                self.logger.info(
-                    f"Reached maximum iteration limit without achieving desired coverage. Current Coverage: {round(self.test_gen.current_coverage * 100, 2)}%"
-                )
+            if self.test_gen.current_coverage < (
+                self.test_gen.desired_coverage / 100
+            ):
+                self.test_gen.run_coverage()
 
-            ReportGenerator.generate_report(
-                test_results_list, self.args.report_filepath
-            )
-        else:
+        if self.test_gen.current_coverage >= (self.test_gen.desired_coverage / 100):
             self.logger.info(
-                f"Prompt only option requested. Skipping call to LLM. Prompt can be found at: {self.args.prompt_only}"
+                f"Reached above target coverage of {self.test_gen.desired_coverage}% (Current Coverage: {round(self.test_gen.current_coverage * 100, 2)}%) in {iteration_count} iterations."
             )
+        elif iteration_count == self.args.max_iterations:
+            failure_message = f"Reached maximum iteration limit without achieving desired coverage. Current Coverage: {round(self.test_gen.current_coverage * 100, 2)}%"
+            if self.args.strict_coverage:
+                # User requested strict coverage (similar to "--cov-fail-under in pytest-cov"). Fail with exist code 2.
+                self.logger.error(failure_message)
+                sys.exit(2)
+            else:
+                self.logger.info(failure_message)
+
+        ReportGenerator.generate_report(
+            test_results_list, self.args.report_filepath
+        )
 
         if 'WANDB_API_KEY' in os.environ:
             wandb.finish()
diff --git a/cover_agent/main.py b/cover_agent/main.py
@@ -81,9 +81,9 @@ def parse_args():
         help="The API url to use for Ollama or Hugging Face. Default: %(default)s.",
     )
     parser.add_argument(
-        "--prompt-only",
+        "--strict-coverage",
         action="store_true",
-        help="Generate the prompt only (i.e. do not run the test or call to the LLM). Default: False.",
+        help="If set, Cover-Agent will return a non-zero exit code if the desired code coverage is not achieved. Default: False.",
     )
 
     return parser.parse_args()

diff --git a/cover_agent/version.txt b/cover_agent/version.txt
@@ -1 +1 @@
-0.1.41
+0.1.42
diff --git a/tests_integration/test_with_docker.sh b/tests_integration/test_with_docker.sh
@@ -76,7 +76,8 @@ COMMAND="/usr/local/bin/cover-agent \
   --test-command \"$TEST_COMMAND\" \
   --coverage-type \"cobertura\" \
   --desired-coverage 70 \
-  --max-iterations 1"
+  --max-iterations 1 \
+  --strict-coverage"
 
 if [ -n "$MODEL" ]; then
   COMMAND="$COMMAND --model \"$MODEL\""