feature/298-flag-disable-file-generation (#324)

* Feature-298: Flag to Disable File Generation

* Feature-298: Flag to Disable File Generation

* Rename flag to suppress-logs

* Fix bug in CoverAgent

* Add files creation suppression logic

* Improve logger usage in classes

* Refactor CustomLogger

* Fix failing test test_main_test_file_not_found

* Extend tests to check if suppress-logs passed or not

* Add file logger test for CustomLogger

* Add log level handling and tests

* Remove get log_level from env var

* Rename suppress-logs to suppress-log-files

* Apply review suggestion for log suppression description

* Improve tests integration scripts

* Remove .coveragerc file due to side effects

---------

Co-authored-by: akashjainwork <[email protected]>

Author: andsimakov

.coveragerc

@@ -215,6 +215,8 @@ A few debug files will be outputted locally within the repository (that are part
   * `stdout`
   * Generated test
 
+You can suppress logs using the `--suppress-log-files` flag. This prevents the creation of the `run.log`, `test_results.html`, and the test results `db` files.
+
 ### Additional logging
 If you set an environment variable `WANDB_API_KEY`, the prompts, responses, and additional information will be logged to [Weights and Biases](https://wandb.ai/).
 

README.md

@@ -42,7 +42,11 @@ def __init__(self, args, agent_completion: AgentCompletionABC=None, logger: Opti
             FileNotFoundError: If required source files or directories are not found.
         """
         self.args = args
-        self.logger = logger or CustomLogger.get_logger(__name__)
+        self.generate_log_files = not args.suppress_log_files
+        # Initialize logger with file generation flag
+        self.logger = logger or CustomLogger.get_logger(__name__, generate_log_files=self.generate_log_files)
+        if args.suppress_log_files:
+            self.logger.info("Suppressed all generated log files.")
 
         self._validate_paths()
         self._duplicate_test_file()
@@ -52,7 +56,9 @@ def __init__(self, args, agent_completion: AgentCompletionABC=None, logger: Opti
             self.agent_completion = agent_completion
         else:
             self.ai_caller = self._initialize_ai_caller()
-            self.agent_completion = DefaultAgentCompletion(caller=self.ai_caller)
+            self.agent_completion = DefaultAgentCompletion(
+            caller=self.ai_caller, generate_log_files=self.generate_log_files
+            )
 
         # Modify test command for a single test execution if needed
         test_command = args.test_command
@@ -103,6 +109,7 @@ def __init__(self, args, agent_completion: AgentCompletionABC=None, logger: Opti
             llm_model=args.model,
             use_report_coverage_feature_flag=args.use_report_coverage_feature_flag,
             agent_completion=self.agent_completion,
+            generate_log_files=self.generate_log_files,
         )
 
         # Initialize test validator with configuration
@@ -124,6 +131,7 @@ def __init__(self, args, agent_completion: AgentCompletionABC=None, logger: Opti
             num_attempts=args.run_tests_multiple_times,
             agent_completion=self.agent_completion,
             max_run_time=args.max_run_time,
+            generate_log_files=self.generate_log_files,
         )
 
     def _initialize_ai_caller(self):
@@ -196,9 +204,9 @@ def _validate_paths(self):
         if not self.args.log_db_path:
             self.args.log_db_path = "cover_agent_unit_test_runs.db"
         # Connect to the test DB
-        self.test_db = UnitTestDB(
-            db_connection_string=f"sqlite:///{self.args.log_db_path}"
-        )
+
+        if self.generate_log_files:
+            self.test_db = UnitTestDB(db_connection_string=f"sqlite:///{self.args.log_db_path}")
 
     def _duplicate_test_file(self):
         """
@@ -263,13 +271,23 @@ def generate_and_validate_tests(self, failed_test_runs, language, test_framework
             ]
 
             # Insert results into database
-            for result in test_results:
-                result["prompt"] = self.test_gen.prompt
-                self.test_db.insert_attempt(result)
+            if self.has_test_db():
+                for result in test_results:
+                    result["prompt"] = self.test_gen.prompt
+                    self.test_db.insert_attempt(result)
 
         except AttributeError as e:
             self.logger.error(f"Failed to validate the tests within {generated_tests_dict}. Error: {e}")
 
+    def has_test_db(self) -> bool:
+        """
+        Check if the test database is initialized.
+
+        Returns:
+            bool: True if the test database is initialized, False otherwise.
+        """
+        return hasattr(self, "test_db") and self.test_db is not None
+
     def check_iteration_progress(self):
         """
         Evaluate current progress towards coverage goals.
@@ -320,8 +338,11 @@ def finalize_test_generation(self, iteration_count):
             f"Total number of output tokens used for LLM model {self.args.model}: {self.test_gen.total_output_token_count + self.test_validator.total_output_token_count}"
         )
 
-        # Generate report and cleanup
-        self.test_db.dump_to_report(self.args.report_filepath)
+        # Only generate report if file generation is enabled
+        if self.generate_log_files:
+            # Generate report and cleanup
+            self.test_db.dump_to_report(self.args.report_filepath)
+
         if "WANDB_API_KEY" in os.environ:
             wandb.finish()
 

cover_agent/CoverAgent.py

@@ -1,5 +1,5 @@
 from cover_agent.CustomLogger import CustomLogger
-from typing import Literal, Tuple, Union, List
+from typing import Literal, Tuple, Union, List, Optional
 import csv
 import json
 import os
@@ -15,6 +15,8 @@ def __init__(
         coverage_type: Literal["cobertura", "lcov", "jacoco"],
         use_report_coverage_feature_flag: bool = False,
         diff_coverage_report_path: str = None,
+        logger: Optional[CustomLogger] = None,
+        generate_log_files: bool = True,
     ):
         """
         Initializes a CoverageProcessor object.
@@ -23,6 +25,8 @@ def __init__(
             file_path (str): The path to the coverage report file.
             src_file_path (str): The fully qualified path of the file for which coverage data is being processed.
             coverage_type (Literal["cobertura", "lcov"]): The type of coverage report being processed.
+            logger (CustomLogger): The logger object for logging messages.
+            generate_log_files (bool): Whether or not to generate logs.
 
         Attributes:
             file_path (str): The path to the coverage report file.
@@ -36,7 +40,7 @@ def __init__(
         self.file_path = file_path
         self.src_file_path = src_file_path
         self.coverage_type = coverage_type
-        self.logger = CustomLogger.get_logger(__name__)
+        self.logger = logger or CustomLogger.get_logger(__name__, generate_log_files=generate_log_files)
         self.use_report_coverage_feature_flag = use_report_coverage_feature_flag
         self.diff_coverage_report_path = diff_coverage_report_path
 

cover_agent/CoverageProcessor.py

@@ -1,21 +1,26 @@
 import logging
-import os
 
 
 class CustomLogger:
+
     @classmethod
-    def get_logger(cls, name):
+    def get_logger(cls, name, generate_log_files=True, file_level=logging.INFO, console_level=logging.INFO):
         """
         Return a logger object with specified name.
 
         Parameters:
             name (str): The name of the logger.
+            generate_log_files (bool): Whether to generate log files.
+            file_level (int): The log level to use.
+            console_level (int): The log level to use.
 
         Returns:
             logging.Logger: The logger object.
 
         Note:
-            This method sets up the logger to handle all messages of DEBUG level and above. It adds a file handler to write log messages to a file specified by 'log_file_path' and a stream handler to output log messages to the console. The log file is overwritten on each run.
+            This method sets up the logger to handle all messages of DEBUG level and above.
+            It adds a file handler to write log messages to a file specified by 'log_file_path' and a stream handler
+            to output log messages to the console. The log file is overwritten on each run.
 
         Example:
             logger = CustomLogger.get_logger('my_logger')
@@ -26,35 +31,27 @@ def get_logger(cls, name):
             logger.critical('This is a critical message')
         """
         logger = logging.getLogger(name)
-        logger.setLevel(
-            logging.DEBUG
-        )  # Set the logger to handle all messages of DEBUG level and above
+        logger.setLevel(logging.DEBUG)
 
         # Specify the log file path
         log_file_path = "run.log"
 
         # Check if handlers are already set up to avoid adding them multiple times
         if not logger.handlers:
-            # File handler for writing to a file
-            file_handler = logging.FileHandler(
-                log_file_path, mode="w"
-            )  # Use 'w' to overwrite the log file on each run
-            file_handler.setLevel(logging.INFO)
-            file_formatter = logging.Formatter(
-                "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-            )
-            file_handler.setFormatter(file_formatter)
-            logger.addHandler(file_handler)
+            formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+
+            # Only add file handler if file generation is enabled
+            if generate_log_files:
+                # File handler for writing to a file. Use 'w' to overwrite the log file on each run
+                file_handler = logging.FileHandler(log_file_path, mode="w")
+                file_handler.setLevel(file_level)
+                file_handler.setFormatter(formatter)
+                logger.addHandler(file_handler)
 
             # Stream handler for output to the console
             stream_handler = logging.StreamHandler()
-            stream_handler.setLevel(
-                logging.INFO
-            )  # Set this to DEBUG if you want to see DEBUG messages in the console
-            stream_formatter = logging.Formatter(
-                "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-            )
-            stream_handler.setFormatter(stream_formatter)
+            stream_handler.setLevel(console_level)
+            stream_handler.setFormatter(formatter)
             logger.addHandler(stream_handler)
 
             # Prevent log messages from being propagated to the root logger

cover_agent/CustomLogger.py

@@ -5,7 +5,7 @@
 from cover_agent.utils import load_yaml
 
 from jinja2 import Environment, StrictUndefined
-from typing import Tuple
+from typing import Optional, Tuple
 
 
 class DefaultAgentCompletion(AgentCompletionABC):
@@ -16,15 +16,17 @@ class DefaultAgentCompletion(AgentCompletionABC):
     to get the response.
     """
 
-    def __init__(self, caller: AICaller):
+    def __init__(self, caller: AICaller, logger: Optional[CustomLogger]=None, generate_log_files: bool=True):
         """
         Initializes the DefaultAgentCompletion.
 
         Args:
             caller (AICaller): A class responsible for sending the prompt to an AI model and returning the response.
+            logger (CustomLogger, optional): The logger object for logging messages.
+            generate_log_files (bool, optional): Whether or not to generate logs.
         """
         self.caller = caller
-        self.logger = CustomLogger.get_logger(__name__)
+        self.logger = logger or CustomLogger.get_logger(__name__, generate_log_files=generate_log_files)
 
     def _build_prompt(self, file: str, **kwargs) -> dict:
         """

cover_agent/DefaultAgentCompletion.py

@@ -1,12 +1,15 @@
 import json
 import os
 
+from typing import Optional
+
 from cover_agent.CustomLogger import CustomLogger
 from cover_agent.FilePreprocessor import FilePreprocessor
 from cover_agent.AgentCompletionABC import AgentCompletionABC
 from cover_agent.settings.config_loader import get_settings
 from cover_agent.utils import load_yaml
 
+
 MAX_TESTS_PER_RUN = 4
 
 
@@ -25,6 +28,8 @@ def __init__(
         additional_instructions: str = "",
         use_report_coverage_feature_flag: bool = False,
         project_root: str = "",
+        logger: Optional[CustomLogger]=None,
+        generate_log_files: bool=True,
     ):
         """
         Initialize the UnitTestGenerator class with the provided parameters.
@@ -45,6 +50,8 @@ def __init__(
             use_report_coverage_feature_flag (bool, optional): Setting this to True considers the coverage of all the files in the coverage report.
                                                                This means we consider a test as good if it increases coverage for a different
                                                                file other than the source file. Defaults to False.
+            logger (CustomLogger, optional): The logger object for logging messages.
+            generate_log_files (bool): Whether or not to generate logs.
 
         Returns:
             None
@@ -64,9 +71,10 @@ def __init__(
         self.last_coverage_percentages = {}
         self.llm_model = llm_model
         self.agent_completion = agent_completion
+        self.generate_log_files = generate_log_files
 
         # Get the logger instance from CustomLogger
-        self.logger = CustomLogger.get_logger(__name__)
+        self.logger = logger or CustomLogger.get_logger(__name__, generate_log_files=self.generate_log_files)
 
         # States to maintain within this class
         self.preprocessor = FilePreprocessor(self.test_file_path)

cover_agent/UnitTestGenerator.py

@@ -1,11 +1,12 @@
-from wandb.sdk.data_types.trace_tree import Trace
 import datetime
 import json
 import logging
 import os
-import re
+
+from typing import Optional
 
 from diff_cover.diff_cover_tool import main as diff_cover_main
+from wandb.sdk.data_types.trace_tree import Trace
 
 from cover_agent.AgentCompletionABC import AgentCompletionABC
 from cover_agent.CoverageProcessor import CoverageProcessor
@@ -36,6 +37,8 @@ def __init__(
         diff_coverage: bool = False,
         comparison_branch: str = "main",
         num_attempts: int = 1,
+        logger: Optional[CustomLogger]=None,
+        generate_log_files: bool=True,
     ):
         """
         Initialize the UnitTestValidator class with the provided parameters.
@@ -57,6 +60,8 @@ def __init__(
             use_report_coverage_feature_flag (bool, optional): Setting this to True considers the coverage of all the files in the coverage report.
                                                                This means we consider a test as good if it increases coverage for a different
                                                                file other than the source file. Defaults to False.
+            logger (CustomLogger, optional): The logger object for logging messages.
+            generate_log_files (bool): Whether or not to generate logs.
 
         Returns:
             None
@@ -84,9 +89,10 @@ def __init__(
         self.num_attempts = num_attempts
         self.agent_completion = agent_completion
         self.max_run_time = max_run_time
+        self.generate_log_files = generate_log_files
 
         # Get the logger instance from CustomLogger
-        self.logger = CustomLogger.get_logger(__name__)
+        self.logger = logger or CustomLogger.get_logger(__name__, generate_log_files=self.generate_log_files)
 
         # Override covertype to be 'diff' if diff_coverage is enabled
         if self.diff_coverage:
@@ -120,6 +126,7 @@ def __init__(
             coverage_type=self.coverage_type,
             use_report_coverage_feature_flag=self.use_report_coverage_feature_flag,
             diff_coverage_report_path=self.diff_cover_report_path,
+            generate_log_files=self.generate_log_files,
         )
 
     def get_coverage(self):

cover_agent/UnitTestValidator.py

@@ -116,6 +116,12 @@ def parse_args():
         default="main",
         help="The branch to compare against when using --diff-coverage. Default: %(default)s.",
     )
+    parser.add_argument(
+        "--suppress-log-files",
+        action="store_true",
+        default=False,
+        help="Suppress all generated log files (HTML, logs, DB files).",
+    )
 
     # Create mutually exclusive group
     group = parser.add_mutually_exclusive_group()

cover_agent/main.py

@@ -337,6 +337,12 @@ def parse_args_full_repo():
         action="store_true",
         help="Enable record mode for LLM responses. Default: False.",
     )
+    parser.add_argument(
+        "--suppress-log-files",
+        action="store_true",
+        default=False,
+        help="Suppress all generated log files (HTML, logs, DB files).",
+    )
     return parser.parse_args()