unit-tests-generator 0.1.0__py3-none-any.whl

unit_tests_generator-0.1.0.dist-info/METADATA

@@ -0,0 +1,236 @@
+Metadata-Version: 2.4
+Name: unit-tests-generator
+Version: 0.1.0
+Summary: Open source project to build a unit tests generator in Python.
+Author: Arnau Fabregat
+License: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Operating System :: OS Independent
+Classifier: License :: OSI Approved :: MIT License
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: crewai[google-genai]>=1.10.1
+Requires-Dist: litellm>=1.82.1
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: networkx>=3.6.1
+Requires-Dist: pytest>=9.0.2
+Requires-Dist: pytest-cov>=7.0.0
+Requires-Dist: typer>=0.23.1
+Dynamic: license-file
+
+# 🛡️ Unit-Tests-Generator
+*Autonomous, RAG-Driven Test Engineering for Python.*
+
+**Unit-Tests-Generator** is an intelligent agentic tool that automates the entire `pytest` lifecycle. Unlike static template generators, it leverages an in-memory knowledge graph of your codebase and RAG to produce context-aware, executable, and validated test suites.
+
+The system operates through a multi-stage pipeline:
+1. **Codebase Graphing**: Maps your repository into an in-memory graph to understand cross-file dependencies and type definitions.
+
+2. **Contextual RAG**: Retrieves relevant context for the LLM, ensuring the generated tests understand your project's unique patterns.
+
+3. **Inference & Guardrails**: Generates test cases via LLM API, filtered through strict structural and security guardrails.
+
+4. **The "Sandbox" Validation**: Automatically executes the generated tests.
+
+5. **Persistence**: Only tests that pass execution and meet coverage criteria are committed to your `/tests` directory.
+
+## Quick Start
+```bash
+# Install the package
+pip install unit-tests-generator
+
+# Generate tests using source and output paths
+utgen -s src/ -t tests/
+```
+
+### Setting up your LLM
+Required environment variables:
+```bash
+# .env file
+LLM_OPENROUTER_MODEL = ""
+LLM_OPENROUTER_API_KEY = ""
+```
+- Powered by CrewAI and OpenRouter: https://docs.crewai.com/en/concepts/llms#open-router.
+- OpenRouter models: https://openrouter.ai/models.
+
+## Table of Contents
+1. [Usage](#usage)
+2. [Examples](#examples)
+3. [Graph structure](#graph-structure)
+4. [Code Quality & Documentation](#code-quality--documentation)
+    - [Pre-commit Hooks](#pre-commit-hooks)
+    - [Unit Testing](#unit-testing)
+5. [Virtual Environment](#virtual-environment)
+    - [Create a new virtualenv with the project's dependencies](#create-a-new-virtualenv-with-the-projects-dependencies)
+    - [Checking if the project's virtual environment is active](#checking-if-the-projects-virtual-environment-is-active)
+    - [Updating the project's dependencies](#updating-the-projects-dependencies)
+6. [LICENSE](#license)
+7. [TODO](#todo)
+
+## Usage
+`utgen` is designed to be simple and terminal-first. Once installed, you can generate tests by pointing the tool to your source code and specifying an output directory.
+
+### CLI Arguments
+| Flag | Shortcut | Description | Required | Default |
+| :--- | :--- | :--- | :--- | :--- |
+| `--src_path` | `-s` | Path to the directory containing your source code. | **Yes** | — |
+| `--test_path` | `-t` | Path to the directory where generated tests will be saved. | **Yes** | — |
+| `--graph_path` | `-g` | Path to the RAG-Graph data (enables advanced context). | No | `None` |
+| `--help` | — | Show the help message and exit. | No | — |
+
+### Basic Command
+Run a standard generation without a RAG-Graph:
+```bash
+utgen -s src/ -t tests/
+```
+### Full Command
+Enable RAG-Graph context for more accurate, context-aware test generation:
+```bash
+utgen --src_path src/ --test_path tests/ --graph_path data/
+```
+
+## Examples
+Upload `repo.graphml` to https://lite.gephi.org/.
+
+- This repository:
+    - Code coverage:
+- Example for repository https://github.com/ArnauFabregat/probability_estimation
+![Diagram](https://raw.githubusercontent.com/ArnauFabregat/unit-tests-generator/main/docs/probability-estimation-repo-graph.png)
+    - Code coverage:
+- Example for repository
+    - Code coverage
+
+## Graph structure
+To extract relevant context, the graph needs:
+
+### 🟢 Nodes for:
+- Files
+- Classes
+- Functions
+- Methods
+- Nested functions
+
+Fields:
+- **id**: canonical ID (e.g., file::...::class::Calibrator / ...::method::Calibrator.fit)
+- **type**: "file" | "class" | "method" | "function" | "nested_function"
+- **name**: symbol name ("Calibrator", "fit", etc.)
+- **file**: file path
+- **signature**: normalized function/method signature "def fit(self, probs, y) -> None"
+- **docstring**: (trimmed) docstring
+- **source**: actual source code of the node
+
+### ➡️ Edges for:
+- **defines** (file → symbol, function → nested function)
+- **has_method** (class → method)
+- **calls** (function/method → called symbol)
+- **references** (function/method → referenced symbol)
+
+Fields:
+- **src**: source node id
+- **dst**: destination node id
+- **rel**: "defines" | "has_method" | "references" | "calls"
+
+### 🧠 Prompt template
+TBD
+
+## Code Quality & Documentation
+### Pre-commit Hooks
+---
+This project uses [pre-commit](https://pre-commit.com/) hooks to enforce code quality standards automatically before each commit. The following hooks are configured:
+
+- **Formatting & File Integrity**: `trailing-whitespace`, `end-of-file-fixer`, `check-yaml`, `check-toml`
+- **Code Linting & Formatting**: `ruff-check`, `ruff-format`
+- **Type Checking**: `mypy`
+
+Pre-commit hooks are automatically installed during virtual environment setup (`uv sync`).
+- To run them for modified (staged) files:
+    ```bash
+    uv run pre-commit run
+    ```
+- To run them for the entire repository:
+    ```bash
+    uv run pre-commit run --all-files
+    ```
+
+### Unit Testing
+---
+Unit tests ensure code reliability and prevent regressions. Tests are written using pytest and should cover critical functionality.
+
+To run all tests:
+```bash
+uv run pytest
+```
+
+To run tests with coverage:
+```bash
+uv run pytest --cov
+```
+
+## Virtual Environment
+### Create a new virtualenv with the project's dependencies
+---
+Install the project's virtual environment and set it as your project's Python interpreter.
+This will also install the project's current dependencies.
+
+Open a terminal in VSCode, then execute the following commands:
+
+1. Install [UV: Python package and project manager](https://docs.astral.sh/uv/getting-started/installation/):
+    * On Mac OSX / Linux: `curl -LsSf https://astral.sh/uv/install.sh | sh`
+    * On Windows [In Powershell]: `powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex"`
+
+2. [optional] To create virtual environment from scratch with `uv`: [Working on projects](https://docs.astral.sh/uv/guides/projects/)
+
+3. If the environment already exists, install the virtual environment. It will be installed in the project's root, under a new directory named `.venv`:
+    * `uv sync`
+    * `uv sync --group dev --group dashboard --group test` to install all the dependency groups
+
+4. Activate the new virtual environment:
+    * On Mac OSX / Linux: `source .venv/bin/activate`
+    * On Windows [In Powershell]: `.venv\Scripts\activate`
+
+5. Configure / install pre-commit hooks:
+    * [Pre-commit](https://pre-commit.com/) is a tool that helps us keep the repository complying with certain formatting and style standards, using the `hooks` configured in the `.pre-commit-config.yaml` file.
+    * Previously installed with `uv sync`.
+
+### Checking if the project's virtual environment is active
+---
+All commands listed here assume the project's virtual env is active.
+
+To ensure so, execute the following command, and ensure it points to: `{project_root}/.venv/bin/python`:
+* On Mac OSX / Linux: `which python`
+* On Windows / Mac OSX / Linux: `python -c "import sys; import os; print(os.path.abspath(sys.executable))"`
+
+If not active, execute the following to activate:
+* On Mac OSX / Linux: `source .venv/bin/activate`
+* On Windows [In Powershell]: `.venv\Scripts\activate`
+
+Alternatively, you can also run any command using the prefix `uv run` and `uv` will make sure that it uses the virtual env's Python executable.
+
+### Updating the project's dependencies
+---
+#### Adding new dependencies
+---
+In order to avoid potential version conflicts, we should use uv's dependency manager to add new libraries additional to the project's current dependenies.
+Open a terminal in VSCode and execute the following commands:
+
+* `uv add {dependency}` e.g. `uv add pandas`
+
+This command will update the project's files `pyproject.toml` and `uv.lock` automatically, which are the ones ensuring all developers and environments have the exact same dependencies.
+
+#### Updating your virtual env with dependencies recently added or removed from the project
+---
+Open a terminal in VSCode and execute the following command:
+* `uv sync`
+
+## License
+MIT. See [LICENSE](LICENSE) for more information.
+
+## TODO
+- Study the possibility to run the tests in docker isolated environment for safety
+- Add the remaining guardrails: not allow to write files from test_*.py
+- Replace crewai by langgraph
+- Publish to PyPI

unit_tests_generator-0.1.0.dist-info/RECORD

@@ -0,0 +1,23 @@
+unit_tests_generator-0.1.0.dist-info/licenses/LICENSE,sha256=cglefN9tjSeVCrm5HTQ-5lczPVR1D3r_IKT_NJ2R2vY,1070
+utgen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+utgen/constants.py,sha256=qz-bCgeXV3ArpW9QvOCVf2k5zuCONBNeHyk7pjQ5Ogw,26
+utgen/logger.py,sha256=Xh1ZxJnwZUbqwcWYqD2gsf6bEaQcaHqq5jbZuQnQmXc,2605
+utgen/main.py,sha256=uYr5Ae9aCEJUPth7ue_EFHKsq3UCzpMA28P3gvoPPos,1518
+utgen/pipeline.py,sha256=3c8nHfQMKWxjBcpoawBGija8Yy47aFVXWUODF4vhkR4,3858
+utgen/validation.py,sha256=pMLcCv8I1U25pIXo_JS1m_X4Ph0zI0KlcCg9EZ4iru8,2828
+utgen/config/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+utgen/config/params.py,sha256=EacHophIxiT2KzEyAiWynsbecrDjqf3_8-wdoks5C6w,19
+utgen/raggraph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+utgen/raggraph/parser.py,sha256=eCczESMEevtzTJs2Kqev3az6SKhsjszyf2Cc7NJBuY8,8329
+utgen/raggraph/utils.py,sha256=VwXqXMvC7EtZEdOwwJn0wHeQRDPICPX7wWgG2OM8Sjo,5854
+utgen/raggraph/walker.py,sha256=a90EsYdGV8dFiKdpHp7T3cvVlbo7EarQ1X18xExgG8Y,2591
+utgen/test_generation_crew/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+utgen/test_generation_crew/crew.py,sha256=YPtZIplumHMylc4KptBC0ZBUgzW-LZWAfwSJhQwxjqk,2992
+utgen/test_generation_crew/guardrails.py,sha256=D94ZL3ocfJvEYxejCd8LYJOZmQB8wndWUd6EQj8Ube0,3805
+utgen/test_generation_crew/llm_config.py,sha256=rsetwm8PCVljM4sQSF1e1s_KmFssrMSclwn6nMi4odk,749
+utgen/test_generation_crew/schemas.py,sha256=oBK9sIdrPLTMPTMdTRy3HFzKsjlFHf2csFUhcclNGkQ,855
+unit_tests_generator-0.1.0.dist-info/METADATA,sha256=lFMX3jZHMmZ5L9F9XKjXtKhQx7X4krMcATDiKT-EB6c,9196
+unit_tests_generator-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+unit_tests_generator-0.1.0.dist-info/entry_points.txt,sha256=j0RKTVAFIdePsOTKG8oQF4Mit_bGqa0UySKvaiaK2L0,41
+unit_tests_generator-0.1.0.dist-info/top_level.txt,sha256=WUMK5MkWCv5d6k8BpwWCgiF02sKTDBpb-JeHvFD-Wiw,6
+unit_tests_generator-0.1.0.dist-info/RECORD,,

unit_tests_generator-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

unit_tests_generator-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+utgen = utgen.main:app

unit_tests_generator-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 ArnauFabregat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

unit_tests_generator-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+utgen

utgen/config/params.py

@@ -0,0 +1 @@
+DEBUG_LOGS = False

utgen/constants.py

@@ -0,0 +1 @@
+GUARDRAIL_MAX_RETRIES = 3

utgen/logger.py

@@ -0,0 +1,87 @@
+"""
+Central logger configuration for the application.
+
+This module sets up Loguru-based logging, featuring split streams for
+standard output and errors, log rotation, and the ability to silence
+noisy third-party library logs.
+"""
+
+import logging
+import sys
+
+from loguru import logger
+
+from utgen.config.params import DEBUG_LOGS
+
+# List of library names (e.g., 'chromadb', 'httpx') to silence
+DEPENDENCIES_WITH_LOGGING: list[str] = []
+
+LOG_FORMAT = (
+    "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
+    "<level>{level: <8}</level> | "
+    "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - "
+    "<level>{message}</level>"
+)
+
+
+def disable_dependency_loggers(dependencies: list[str]) -> None:
+    """
+    Disables logging for specific third-party dependencies.
+
+    Prevents external library logs from cluttering the application logs
+    by setting their 'disabled' flag to True and stopping propagation.
+
+    Parameters
+    ----------
+    dependencies : list[str]
+        A list of library names as strings to be disabled.
+    """
+    for name in dependencies:
+        logging.getLogger(name).disabled = True
+        logging.getLogger(name).propagate = False
+
+
+def setup_logger(debug: bool | None = None) -> None:
+    """
+    Configures and initializes the Loguru logger handlers.
+
+    Sets up a three-tier logging strategy:
+    1. Standard Output: Handles DEBUG/INFO levels (Severity < 30).
+    2. Standard Error: Handles WARNING and higher (Severity >= 30).
+    3. File: Persistent storage with rotation and compression.
+
+    Parameters
+    ----------
+    debug : bool, optional
+        If True, the base log level is set to 'DEBUG'. If False, it
+        defaults to 'INFO'. If None, it uses the logic defined
+        within the function.
+    """
+    # Remove default Loguru handler
+    logger.remove()
+
+    # Determine the "Floor" level (DEBUG or INFO)
+    base_level = "DEBUG" if debug else "INFO"
+
+    # 1. Console: Standard Output (DEBUG/INFO)
+    # Uses a filter to ensure higher severity levels don't duplicate to stdout
+    logger.add(
+        sys.stdout,
+        format=LOG_FORMAT,
+        level=base_level,
+        filter=lambda r: r["level"].no < 30,  # Stop before WARNING
+    )
+
+    # 2. Console: Standard Error (WARNING+)
+    # This ensures critical issues are highlighted in the error stream
+    logger.add(
+        sys.stderr,
+        format="\n" + LOG_FORMAT,
+        level="WARNING",
+    )
+
+
+# --- INITIALIZATION ---
+# Run once when module is imported to apply settings globally
+disable_dependency_loggers(DEPENDENCIES_WITH_LOGGING)
+setup_logger(debug=DEBUG_LOGS)

utgen/main.py

@@ -0,0 +1,43 @@
+import subprocess
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+from utgen.logger import logger
+from utgen.pipeline import pipeline
+
+app = typer.Typer(help="UTGen: Unit test generation with RAG-Graph and LLMs")
+
+
+@app.callback(invoke_without_command=True)
+def run(
+    src_path: Annotated[Path, typer.Option("--src_path", "-s", help="Path to source")],
+    test_path: Annotated[Path, typer.Option("--test_path", "-t", help="Path to tests")],
+    graph_path: Annotated[Path | None, typer.Option("--graph_path", "-g", help="Path to Graph")] = None,
+) -> None:
+    """
+    Generate unit tests based on your source code and RAG-Graph.
+    """
+    logger.info("Initializing utgen...")
+    logger.info(f"Source: {src_path} | Tests: {test_path} | Graph: {graph_path}")
+
+    pipeline(
+        source_code_dir=str(src_path),
+        tests_output_dir=str(test_path),
+        save_graph_path=str(graph_path) if graph_path else "",
+    )
+
+    logger.info("Running coverage report...")
+    # We run pytest-cov on the newly generated test directory
+    # --cov points to the source code we want to measure
+    try:
+        subprocess.run(["pytest", str(test_path), f"--cov={src_path}", "--cov-report=term-missing"], check=True)
+    except subprocess.CalledProcessError:
+        logger.warning("Some tests failed or coverage couldn't be calculated.")
+    except FileNotFoundError:
+        logger.error("`pytest` not found. Make sure it's installed in your environment.")
+
+
+if __name__ == "__main__":
+    app()

utgen/pipeline.py

@@ -0,0 +1,84 @@
+import json
+from collections import defaultdict
+from pathlib import Path
+
+from utgen.logger import logger
+from utgen.raggraph.utils import get_node_context
+from utgen.raggraph.walker import build_graph_from_directory
+from utgen.test_generation_crew.crew import TestGenerationCrew
+from utgen.validation import save_and_clean_tests, validate_individual_test
+
+
+def pipeline(source_code_dir: str, tests_output_dir: str, save_graph_path: str = "") -> None:
+    """
+    Main function to orchestrate the test generation process.
+    Args:
+        source_code_dir (str): Directory containing the source code to analyze.
+        tests_output_dir (str): Directory where generated tests will be saved.
+        save_graph_path (str): Path to save the graph representation.
+    """
+    logger.info("Creating graph from source code...")
+    g = build_graph_from_directory(code_path=source_code_dir, save_graph_path=save_graph_path)
+    logger.info(f"Graph built with {g.number_of_nodes()} nodes and {g.number_of_edges()} edges.")
+
+    logger.info("Started test generation process...")
+    # Define defaultdict of dicts
+    tests_results: defaultdict[str, dict[str, dict]] = defaultdict(dict)
+
+    # TODO: afegir guardrails que falten
+    test_generator = TestGenerationCrew(guardrail_max_retries=5, verbose=False)
+
+    for node_id, data in list(g.nodes(data=True))[20:30]:
+        if data["type"] in ["function", "method"]:
+            logger.info(f"Generating tests for node: {node_id}")
+            try:
+                # Get context
+                context = get_node_context(g=g, node_id=node_id)
+                inputs = {"graph_context": context}
+
+                # Generate tests
+                response = test_generator.crew().kickoff(inputs=inputs)
+
+                # Convert string to dictionary
+                response_dict = json.loads(response.raw)
+
+                # Store results
+                p = Path(data["file"])
+                new_filename = f"test_{p.stem}{p.suffix}"
+                save_path = (p.parent / new_filename).as_posix()
+                tests_results[save_path][node_id] = response_dict["tests"]
+
+            except Exception:
+                # This catches guardrail retries exceeded, JSON parsing errors, etc.
+                logger.error(f"Failed to generate tests for {node_id} after max retries.")
+                # Use 'continue' to skip the rest of this iteration and move to the next node
+                continue
+    logger.info("Test generation process completed.")
+
+    logger.info("Validating and saving generated tests...")
+    for save_path, nodes in tests_results.items():
+        accepted_tests: list[tuple[str, str]] = []
+        for node_id, tests in nodes.items():
+            base_import = (
+                "from "
+                + node_id.split("::")[0][:-3].replace("/", ".")
+                + " import "
+                + node_id.split("::")[-1].split(".")[0]
+            )
+            for name, values in tests.items():
+                imports, code = values["imports"], values["code"]
+                if base_import not in imports:
+                    logger.debug(f"Added missing import `{base_import}` for test `{name}`.")
+                    imports.append(base_import)
+                imports = "\n".join(imports)
+
+                if validate_individual_test(import_code=imports, test_code=code):
+                    logger.debug(f"Test `{name}` accepted during validation.")
+                    accepted_tests.append((imports, code))
+                else:
+                    logger.debug(f"Test `{name}` rejected during validation.")
+
+        logger.debug(f"Saving cleaned tests for `{save_path}`...")
+        save_and_clean_tests(valid_tests=accepted_tests, destination=f"{tests_output_dir}/{save_path}")
+    logger.info("All tests validated and saved successfully.")
+    # TODO: run pytest coverage