Skip to content

Commit 810cc5d

Browse files
committed
Updated Java test generation example
Signed-off-by: Saurabh Sinha <sinha108@gmail.com>
1 parent 475a5ed commit 810cc5d

1 file changed

Lines changed: 82 additions & 79 deletions

File tree

docs/examples/java/notebook/generate_unit_tests.ipynb

Lines changed: 82 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -9,29 +9,21 @@
99
"source": [
1010
"# Using CLDK to generate JUnit tests\n",
1111
"\n",
12-
"In this tutorial, we will use CLDK to generate a JUnit test for all the methods in a Java Application.\n",
12+
"In this tutorial, we will use CLDK to implement a simple unit test generator for Java. You'll explore some of the benefits of using CLDK to perform quick and easy program analysis and build an LLM-based test generator. By the end of this tutorial, you will have implemented such a tool and generated a JUnit test case for a Java application.\n",
1313
"\n",
14-
"By the end of this tutorial, you will have a JUnit test for all the methods in a Java application. You'll be able to explore some of the benefits of using CLDK to perform fast and easy program analysis and build a LLM-based test generator.\n",
14+
"Specifically, you will learn how to perform the following tasks on the application under test to create LLM prompts for test generation:\n",
1515
"\n",
16-
"You will learn how to do the following:\n",
16+
"1. Create a new instance of the CLDK class.\n",
17+
"2. Create an analysis object for the Java application under test.\n",
18+
"3. Iterate over all files in the application.\n",
19+
"4. Iterate over all classes in a file.\n",
20+
"5. Iterate over all methods in a class.\n",
21+
"6. Get the code body of a method.\n",
22+
"7. Get the constructors of a class.\n",
23+
"<!-- 7. Initialize treesitter utils for the class file content.\n",
24+
"8. Sanitize the class for analysis. -->\n",
1725
"\n",
18-
"<ol>\n",
19-
"<li> Create a new instance of the CLDK class.\n",
20-
"<li> Create an analysis object over the Java application.\n",
21-
"<li> Iterate over all the files in the project.\n",
22-
"<li> Iterate over all the classes in the file.\n",
23-
"<li> Iterate over all the methods in the class.\n",
24-
"<li> Get the code body of the method.\n",
25-
"<li> Initialize the treesitter utils for the class file content.\n",
26-
"<li> Sanitize the class for analysis.\n",
27-
"</ol>\n",
28-
"Next, we will write a couple of helper methods to:\n",
29-
"\n",
30-
"<ol>\n",
31-
"<li> Format the instruction for the given focal method and class.\n",
32-
"<li> Prompts the local model on Ollama.\n",
33-
"<li> Use CLDK to go through an application and generate unit test cases for each method.\n",
34-
"</ol>"
26+
"We will write a couple of helper methods to (1) format the LLM instruction for generating test cases for a given focal method (i.e., method under test) and (2) prompt the LLM via Ollama. We will then use CLDK to go through an application and generate unit test cases for a target method."
3527
]
3628
},
3729
{
@@ -45,12 +37,11 @@
4537
"\n",
4638
"Before we get started, let's make sure you have the following installed:\n",
4739
"\n",
48-
"<ol>\n",
49-
"<li> Python 3.11 or later\n",
50-
"<li> Ollama 0.3.4 or later\n",
51-
"<li> Java 11 or later\n",
52-
"</ol>\n",
53-
"We will use ollama to spin up a local granite model that will act as our LLM for this turorial."
40+
"1. Python 3.11 or later (you can use [pyenv](https://github.com/pyenv/pyenv) to install Python)\n",
41+
"2. Java 11 or later (you can use [SDKMAN!](https://sdkman.io) to instal Java)\n",
42+
"3. Ollama 0.3.4 or later (you can get Ollama here: [Ollama download](https://ollama.com/download))\n",
43+
"\n",
44+
"We will use Ollama to spin up a local [Granite code model](https://ollama.com/library/granite-code), which will serve as our LLM for this turorial."
5445
]
5546
},
5647
{
@@ -60,12 +51,28 @@
6051
"collapsed": false
6152
},
6253
"source": [
63-
"### Prerequisite 1: Install ollama\n",
54+
"### Download Granite code model\n",
6455
"\n",
65-
"If you don't have ollama installed, please download and install it from here: [Ollama](https://ollama.com/download).\n",
66-
"Once you have ollama, start the server and make sure it is running. Once ollama is up and running, you can download the latest version of the Granite 8b Instruct model by running the following command:\n",
67-
"There are other granite versions available, but for this tutorial, we will use the Granite 8b Instruct model. You if prefer to use a different version, you can replace `8b-instruct` with any of the other [versions](https://ollama.com/library/granite-code/tags).\n",
68-
"Let's make sure the model is downloaded by running the following command:"
56+
"After starting the Ollama server, please download the latest version of the Granite code 8b-instruct model by running the following command. There are other Granite code models available, but for this tutorial, we will use Granite code 8b-instruct. If you prefer to use a different Granite code model, you can replace `8b-instruct` with the tag of another version (see [Granite code tags](https://ollama.com/library/granite-code/tags))."
57+
]
58+
},
59+
{
60+
"cell_type": "code",
61+
"execution_count": null,
62+
"id": "670f2b23",
63+
"metadata": {},
64+
"outputs": [],
65+
"source": [
66+
"%%bash\n",
67+
"ollama pull granite-code:8b-instruct"
68+
]
69+
},
70+
{
71+
"cell_type": "markdown",
72+
"id": "02d5bbfa",
73+
"metadata": {},
74+
"source": [
75+
" Let's make sure the model is downloaded by running the following command: "
6976
]
7077
},
7178
{
@@ -82,7 +89,7 @@
8289
"outputs": [],
8390
"source": [
8491
"%%bash\n",
85-
"ollama run granite-code:8b-instruct \\\"Write a python function to print 'Hello, World!'"
92+
"ollama run granite-code:8b-instruct \\\"Write a python function to print 'Hello, World!'\\\""
8693
]
8794
},
8895
{
@@ -92,7 +99,7 @@
9299
"collapsed": false
93100
},
94101
"source": [
95-
"### Prerequisite 3: Install ollama Python SDK"
102+
"### Install Ollama Python SDK"
96103
]
97104
},
98105
{
@@ -114,8 +121,8 @@
114121
"collapsed": false
115122
},
116123
"source": [
117-
"### Prerequisite 4: Install CLDK\n",
118-
"CLDK is avaliable on github at github.com/IBM/codellm-devkit.git. You can install it by running the following command:"
124+
"### Install CLDK\n",
125+
"CLDK is avaliable at https://github.com/IBM/codellm-devkit. You can install it by running the following command:"
119126
]
120127
},
121128
{
@@ -138,7 +145,7 @@
138145
},
139146
"source": [
140147
"### Get the sample Java application\n",
141-
"For this tutorial, we will use apache commons cli. You can download the source code to a temporary directory by running the following command:"
148+
"For this tutorial, we will use [Apache Commons CLI](https://github.com/apache/commons-cli) as the Java application under test. You can download the source code to a temporary directory by running the following command:"
142149
]
143150
},
144151
{
@@ -161,7 +168,7 @@
161168
"collapsed": false
162169
},
163170
"source": [
164-
"The project will now be extracted to `/tmp/commons-cli-rel-commons-cli-1.7.0`. We'll remove these files later, so don't worry about the location."
171+
"The project will be extracted to `/tmp/commons-cli-rel-commons-cli-1.7.0`. We'll remove these files later, so don't worry about the location."
165172
]
166173
},
167174
{
@@ -171,22 +178,16 @@
171178
"collapsed": false
172179
},
173180
"source": [
174-
"### Building a JUnit test generator using CLDK and Granite Code Instruct Model\n",
175-
"Now that we have all the prerequisites installed, let's start building a JUnit test generator using CLDK and the Granite Code Instruct Model."
176-
]
177-
},
178-
{
179-
"cell_type": "markdown",
180-
"id": "5856baff4aa64ed7",
181-
"metadata": {
182-
"collapsed": false
183-
},
184-
"source": [
185-
"Generating unit tests for code is a very tedious task and often takes a significant effort from the developers to write good test cases. There are various tools that are available for automated test generation, such as EvoSuite, which uses evolutionary algorithms to generate test cases. However, the test cases that are being generated are not natural and often developers do not prefer to add them to their test suite. Whereas Large Language Models (LLM) being trained with developer-written code it has a better affinity towards generating more natural code--more readable, maintainable code. In this excercise, we will show we can leverage LLMs to generate test cases with the help of CLDK. \n",
181+
"## Building a JUnit test generator using CLDK and Granite Code Instruct Model\n",
182+
"\n",
183+
"Now that we have all the prerequisites installed, let's start building a JUnit test generator using CLDK and the Granite Code Instruct Model.\n",
184+
"\n",
185+
"Generating unit tests for code is a tedious task and developers often have to put in significant effort in writing good test cases. There are various tools available for automated test generation, such as EvoSuite, which uses evolutionary algorithms to generate unit test cases for Java. However, the generated test cases are not natural and often developers do not prefer to add them to their test suites. Large Language Models (LLMs), having been trained with developer-written code, have a better affinity towards generating more natural code---code that is more readable, comprehensible, and maintainable. In this excercise, we will show how we can leverage LLMs to generate test cases with the help of CLDK. \n",
186186
"\n",
187-
"For simplicity, we will cover certain aspects of test generation and provide some context information to LLM for better quality of test cases. In this exercise, we will generate a unit test for a non-private method from a Java class and provide the focal method body and the signature of all the constructors of the class so that LLM can understand how to create an object of the focal class during the setup phase of the tests. Also, we will ask LLMs to generate ```N``` number of test cases, where ```N``` is the cyclomatic complexity of the focal method. The intuition is that one test may not be sufficient for covering fairly complex methods, and a cyclomatic complexity score can provide some guidance towards that. \n",
187+
"For simplicity, we will cover certain aspects of test generation and provide some context information to the LLM to help it create usable test cases. In this exercise, we will generate a unit test for a non-private method from a Java class and provide the focal method body and the signature of all the constructors of the class so that LLM can understand how to create an object of the focal class during the setup phase of the tests.\n",
188+
"<!-- Also, we will ask LLMs to generate ```N``` number of test cases, where ```N``` is the cyclomatic complexity of the focal method. The intuition is that one test may not be sufficient for covering fairly complex methods, and a cyclomatic complexity score can provide some guidance towards that. -->\n",
188189
"\n",
189-
"(Step 1) First, we will import all the necessary libraries"
190+
"(Step 1) First, we will import all the necessary libraries."
190191
]
191192
},
192193
{
@@ -210,7 +211,7 @@
210211
"collapsed": false
211212
},
212213
"source": [
213-
"(Step 2) Second, we will form the prompt for the model, which will include all the constructor signarures, and the body of the focal method."
214+
"(Step 2) Second, we will define a function for creating the LLM prompt, which includes signatures of relevant constructors and the body of the focal method."
214215
]
215216
},
216217
{
@@ -224,7 +225,7 @@
224225
"source": [
225226
"def format_inst(focal_method_body, focal_method, focal_class, constructor_signatures, language):\n",
226227
" \"\"\"\n",
227-
" Format the instruction for the given focal method and class.\n",
228+
" Format the LLM instruction for the given focal method and class.\n",
228229
" \"\"\"\n",
229230
" inst = f\"Question: Can you generate junit tests with @Test annotation for the method `{focal_method}` in the class `{focal_class}` below. Only generate the test and no description.\\n\"\n",
230231
" inst += 'Use the constructor signatures to form the object if the method is not static. Generate the code under ``` code block.'\n",
@@ -246,7 +247,7 @@
246247
"collapsed": false
247248
},
248249
"source": [
249-
"(Step 3) Third, use ollama to call LLM (in case Granite 8b)."
250+
"(Step 3) Third, we define a function to call the LLM (in this case, Granite code 8b-instruct) using Ollama."
250251
]
251252
},
252253
{
@@ -271,7 +272,7 @@
271272
"collapsed": false
272273
},
273274
"source": [
274-
"(Step 4) Fourth, collect all the information needed for each method. In this process, we go through all the classes in the application, and then for each class, we collect the signature of all the constructors. If there is no constructor present, we add the signature of the default constructor. Then, we go through all the non-private methods of the class and formulate the prompt using the constructor and the method information. Finally, we use the prompt to call LLM and get the final output."
275+
"(Step 4) Fourth, we collect the relevant information for the focal method. To do this, we go through all the classes in the application, and for each class, we collect the signatures of its constructors. If a class has no constructors, we add the signature of the default constructor. Then, we go through each non-private method of the class and formulate the prompt using the constructor and the method information. Finally, we use the prompt to call LLM to generate test cases and get the LLM response."
275276
]
276277
},
277278
{
@@ -283,27 +284,28 @@
283284
},
284285
"outputs": [],
285286
"source": [
286-
"# Create a new instance of the CLDK class\n",
287+
"# Create an instance of the CLDK class for Java analysis\n",
287288
"cldk = CLDK(language=\"java\")\n",
288-
"# Create an analysis object over the java application. Provide the application path.\n",
289+
"\n",
290+
"# Create an analysis object for the Java application. Provide the application path.\n",
289291
"analysis = cldk.analysis(project_path=\"/tmp/commons-cli-rel-commons-cli-1.7.0\", analysis_level=AnalysisLevel.symbol_table)\n",
290292
"\n",
291-
"# For simplicity, we run the test generation for a single class and method. One can remove that filter to run this code for the entire application\n",
292-
"qualified_class_name = 'org.apache.commons.cli.GnuParser'\n",
293-
"method_signature = 'flatten(Options, String[], boolean)'\n",
293+
"# For simplicity, we run the test generation for a single focal class and method (this filter can be removed to run this code over the entire application)\n",
294+
"focal_class = \"org.apache.commons.cli.GnuParser\"\n",
295+
"focal_method = \"flatten(Options, String[], boolean)\"\n",
294296
"\n",
295297
"# Go through all the classes in the application\n",
296298
"for class_name in analysis.get_classes():\n",
297299
"\n",
298-
" if class_name == qualified_class_name:\n",
300+
" if class_name == focal_class:\n",
299301
" class_details = analysis.get_class(qualified_class_name=class_name)\n",
300-
" focal_class_name = class_name.split('.')[-1]\n",
302+
" focal_class_name = class_name.split(\".\")[-1]\n",
301303
"\n",
302304
" # Generate test cases for non-interface and non-abstract classes\n",
303-
" if not class_details.is_interface and 'abstract' not in class_details.modifiers:\n",
305+
" if not class_details.is_interface and \"abstract\" not in class_details.modifiers:\n",
304306
" \n",
305307
" # Get all constructor signatures\n",
306-
" constructor_signatures = ''\n",
308+
" constructor_signatures = \"\"\n",
307309
" \n",
308310
" for method in analysis.get_methods_in_class(qualified_class_name=class_name):\n",
309311
" method_details = analysis.get_method(qualified_class_name=class_name, qualified_method_name=method)\n",
@@ -312,35 +314,36 @@
312314
" constructor_signatures += method_details.signature + '\\n'\n",
313315
" \n",
314316
" # If no constructor present, then add the signature of the default constructor\n",
315-
" if constructor_signatures=='':\n",
316-
" constructor_signatures = f'public {focal_class_name}() ' + '{}'\n",
317+
" if constructor_signatures == \"\":\n",
318+
" constructor_signatures = f\"public {focal_class_name}() \" + \"{}\"\n",
317319
" \n",
318320
" # Go through all the methods in the class\n",
319321
" for method in analysis.get_methods_in_class(qualified_class_name=class_name):\n",
320322
" \n",
321-
" if method==method_signature:\n",
323+
" if method == focal_method:\n",
322324
" # Get the method details\n",
323325
" method_details = analysis.get_method(qualified_class_name=class_name, qualified_method_name=method)\n",
324326
" \n",
325327
" # Generate test cases for non-private methods\n",
326-
" if 'private' not in method_details.modifiers and not method_details.is_constructor:\n",
328+
" if \"private\" not in method_details.modifiers and not method_details.is_constructor:\n",
327329
" \n",
328330
" # Gather all the information needed for the prompt, which are focal method body, focal method name, focal class name, and constructor signature\n",
329-
" prompt = format_inst(focal_method_body=method_details.declaration+method_details.code,\n",
330-
" focal_method=method.split('(')[0],\n",
331-
" focal_class=focal_class_name,\n",
332-
" constructor_signatures=constructor_signatures,\n",
333-
" language='Java')\n",
331+
" prompt = format_inst(\n",
332+
" focal_method_body=method_details.declaration+method_details.code,\n",
333+
" focal_method=method.split(\"(\")[0],\n",
334+
" focal_class=focal_class_name,\n",
335+
" constructor_signatures=constructor_signatures,\n",
336+
" language=\"Java\"\n",
337+
" )\n",
334338
" \n",
339+
" # Print the instruction\n",
335340
" print(f\"Instruction:\\n{prompt}\\n\")\n",
336-
" print(f\"Generating test case . . .\\n\")\n",
341+
" print(f\"Generating test case ...\\n\")\n",
337342
" \n",
338343
" # Prompt the local model on Ollama\n",
339-
" llm_output = prompt_ollama(\n",
340-
" message=prompt\n",
341-
" )\n",
344+
" llm_output = prompt_ollama(message=prompt)\n",
342345
" \n",
343-
" # Print the instruction and LLM output\n",
346+
" # Print the LLM output\n",
344347
" print(f\"LLM Output:\\n{llm_output}\")"
345348
]
346349
}
@@ -361,7 +364,7 @@
361364
"name": "python",
362365
"nbconvert_exporter": "python",
363366
"pygments_lexer": "ipython3",
364-
"version": "3.11.4"
367+
"version": "3.11.9"
365368
}
366369
},
367370
"nbformat": 4,

0 commit comments

Comments
 (0)