diff --git a/package-lock.json b/package-lock.json
index 5f94117..d6ab1e4 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "keployio",
-  "version": "1.0.15",
+  "version": "1.0.17",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "keployio",
-      "version": "1.0.15",
+      "version": "1.0.17",
       "dependencies": {
         "acorn-walk": "^8.3.3",
         "axios": "^1.7.4",
diff --git a/src/Utg.ts b/src/Utg.ts
index 365a8d3..cbe1e36 100644
--- a/src/Utg.ts
+++ b/src/Utg.ts
@@ -145,8 +145,9 @@ async function Utg(context: vscode.ExtensionContext , additional_prompts?:string
 
             // Define the test file name by appending 'Test' to the original class name
             const originalClassName = path.basename(sourceFilePath, '.java');
-            const testFileName = `${originalClassName}Test.java`;
+            const testFileName = `${originalClassName}Tests.java`;
             const testFilePath = path.join(fullTestDir, testFileName);
+            const JavaClassName = `${originalClassName}Tests`
 
             testFilePaths.push(testFilePath);
 
@@ -154,7 +155,10 @@ async function Utg(context: vscode.ExtensionContext , additional_prompts?:string
                 vscode.window.showInformationMessage("Test doesn't exist", testFilePath);
 
                 // **Create Test File Content with Proper Package and JUnit Imports**
-                const testFileContent = `package ${packageName};`;
+                const testFileContent = `package ${packageName};\n\n
+                    class ${JavaClassName} {\n
+                }\n`;
+
                 fs.writeFileSync(testFilePath, testFileContent);
                 vscode.window.showInformationMessage(`Created test file: ${testFilePath}`);
                 console.log(`🐰 Created test file with package name: ${testFilePath}`);
diff --git a/src/extension.ts b/src/extension.ts
index de59bd1..5b6cab4 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -122,6 +122,7 @@ class KeployCodeLensProvider implements vscode.CodeLensProvider {
         if (
             fileName.endsWith('.test.js') ||
             fileName.endsWith('.test.ts') ||
+            fileName.endsWith('Tests.java') ||  // Check for Java test file ending
             fileName.endsWith('Test.java') ||  // Check for Java test file ending
             fileName.includes('/Test') ||      // Check for Java test file prefix in the path
             fileName.includes('/test/') ||    // Skip files in a "tests" directory
@@ -296,7 +297,11 @@ async function findTestCasesForFunction(functionName: string, fileExtension: str
                 continue;
             }
 
-            const tree = parser.parse(text);
+            const options: TreeSitter.Options = {
+                bufferSize: 1024 * 1024,
+            };
+
+            const tree = parser.parse(text, undefined, options);  // Parse the document text
             const cursor = tree.walk();
             let found = false;
 
@@ -414,8 +419,12 @@ async function getAllFunctionsInFile(
         console.log('🐰 Unsupported file type:', filePath);
         throw new Error("Unsupported file type");
     }
+    const options: TreeSitter.Options = {
+        bufferSize: 1024 * 1024,
+    };
+
+    const tree = parser.parse(text, undefined, options);  // Parse the document text
 
-    const tree = parser.parse(text);
     const cursor = tree.walk();
     const functionNames: string[] = [];