Unit Test Recommendation Engine

This project recommends unit tests for a piece of code or a short description. It summarizes each test, embeds that summary, and retrieves the closest tests by cosine similarity.

Installation

python -m pip install -r requirements.txt

The first run downloads Salesforce/codet5-base-multi-sum and sentence-transformers/multi-qa-MiniLM-L6-dot-v1. CUDA device 0 is used when CUDA is available. CPU is used when CUDA is not available. The summarizer is loaded with the tokenizer and seq2seq model directly, so it runs on both Transformers 4 and Transformers 5.

Optional extras:

• requirements-dev.txt installs pytest.
• requirements-plot.txt installs the dependencies for plot.py.
• requirements-finetune.txt installs salesforce-codetf for finetune.py. That package provides the codetf module and targets older Python releases, so install it in a separate environment when it conflicts with the runtime packages.

Usage

Text similarity:

python main.py 1 "A test for a function that adds two numbers"

Code similarity:

python main.py 2 "def add_numbers(a, b): return a + b"

The command writes the five closest tests to results.txt. --top-k changes how many tests are returned:

python main.py 1 "A test for a function that adds two numbers" --top-k 10

With no arguments, the program asks for the mode and the query.

Timing is printed in plain text and as one machine-readable line:

TIMING {"get_user_input_time": 0.0, "load_unit_tests_time": 0.1, "model_init_time": 1.2, "similarity_time": 0.05, "total_execution_time": 1.35}

test.py reads that line. Run it from the repository root:

python test.py

How it works

1. Load the summarizer and the embedding model.
2. Read unit tests from unittests/, including nested directories. The folder holds the original Python examples, unit tests copied from OpenTailcat, and one JUnit test from each of 1,010 repositories in the Methods2Test test split. Summaries and embeddings are cached in unittest_embeddings.json. Added, edited, and deleted tests update that cache. The FAISS index is rebuilt from the cache, so every vector maps to a file that still exists.
3. Embed the query. Code queries are summarized first. Text queries are embedded directly.
4. Rank tests by cosine similarity. Embeddings are L2-normalized and searched with an inner-product index, which orders candidates by that same score. Results are written to results.txt.

On those 1,010 Methods2Test methods, a word search over the test source placed the mapped test first more often than this index: Hit@1 0.63 against 0.33, and Hit@5 0.83 against 0.50.

The program reads and writes unittest_embeddings.json only. An older unittest_embeddings.pkl file is ignored, because loading a pickle can execute code stored in the file. Delete that file if it is still present.

Tests

python -m pip install -r requirements-dev.txt
python -m pytest

The tests cover index updates, deletions, cosine ranking, empty indexes, argument validation, and the benchmark parser. They do not download the models.

Contributing

Suggestions, bug reports, and pull requests are welcome.

License

This project is licensed under the Apache License, Version 2.0. See the LICENSE file for details.