Add Rust Support for Unit Test Generation (UTG)

package.json

@@ -57,7 +57,7 @@
     },
     "submenus": [
       {
-        "icon": "$(account)",  
+        "icon": "$(account)",
         "label": "Sign In Options",
         "id": "sign_in_submenu"
       }
@@ -107,7 +107,6 @@
         {
           "command": "keploy.SignInWithMicrosoft",
           "title": "Sign In with Microsoft"
-
         }
       ]
     },
@@ -208,10 +207,11 @@
     "sinon": "^17.0.1",
     "svelte": "^4.2.12",
     "tree-sitter": "^0.21.1",
-    "tree-sitter-java": "^0.21.0",
     "tree-sitter-go": "^0.23.0",
+    "tree-sitter-java": "^0.21.0",
     "tree-sitter-javascript": "^0.21.4",
     "tree-sitter-python": "^0.21.0",
+    "tree-sitter-rust": "^0.23.0",
     "uuid": "^10.0.0",
     "walk": "^2.3.15",
     "yaml": "^2.4.2"

scripts/utg.sh

@@ -51,6 +51,38 @@ if [ "$extension" = "go" ]; then
   export PATH=$PATH:$(go env GOPATH)/bin
 fi
 
+if [ "$extension" = "rs" ]; then
+  echo "Setting up Rust testing environment..."
+
+  # Check if rustup is installed
+  if ! command -v rustup &> /dev/null; then
+    echo "rustup is not installed. Please install Rust toolchain first."
+    exit 1
+  fi
+
+  # Check if cargo-llvm-cov is installed
+  if ! cargo install --list | grep -q "cargo-llvm-cov"; then
+    echo "Installing cargo-llvm-cov..."
+    cargo install cargo-llvm-cov
+  else
+    echo "cargo-llvm-cov is already installed."
+  fi
+
+  # Check if the project has a Cargo.toml file
+  if [ ! -f "Cargo.toml" ]; then
+    echo "No Cargo.toml found. Please ensure you're in a Rust project directory."
+    exit 1
+  fi
+
+  # Install llvm-tools-preview component if not already installed
+  rustup component add llvm-tools-preview
+
+  # Set default test command if none provided
+  if [ -z "$command" ]; then
+    command="cargo llvm-cov --html --output-dir \"$coverageReportPath\""
+  fi
+fi
+
 # Construct the keploy gen command
 if [ "$extension" = "java" ]; then
   keployCommand="keploy gen --source-file-path=\"$sourceFilePath\" \
@@ -61,6 +93,15 @@ if [ "$extension" = "java" ]; then
     --llmBaseUrl \"https://api.keploy.io\" \
     --max-iterations \"5\" \
     --coverageFormat jacoco"
+elif [ "$extension" = "rs" ]; then
+  keployCommand="keploy gen --source-file-path=\"$sourceFilePath\" \
+    --test-file-path=\"$testFilePath\" \
+    --test-command=\"$command\" \
+    --coverage-report-path=\"$coverageReportPath\" \
+    --llmApiVersion \"2024-02-01\" \
+    --llmBaseUrl \"https://api.keploy.io\" \
+    --max-iterations \"5\" \
+    --coverageFormat lcov"
 else
   keployCommand="keploy gen --source-file-path=\"$sourceFilePath\" \
     --test-file-path=\"$testFilePath\" \
@@ -78,5 +119,4 @@ if [ -n "$additional_prompts" ] && [ "$additional_prompts" != " " ]; then
 fi
 
 # Run the keploy command
-# echo "Running: $keployCommand"
-eval $keployCommand
+eval $keployCommand

src/Utg.ts

@@ -150,7 +150,7 @@
             const originalClassName = path.basename(sourceFilePath, '.java');
             const testFileName = `${originalClassName}Tests.java`;
             const testFilePath = path.join(fullTestDir, testFileName);
             const JavaClassName = `${originalClassName}Tests`
 
             testFilePaths.push(testFilePath);
 
@@ -215,6 +215,53 @@
                    // **Set Command and Coverage Report Path for Go**
                    command = `go test -v ./... -coverprofile=coverage.out && gocov convert coverage.out | gocov-xml > coverage.xml`;
                    coverageReportPath = "./coverage.xml";
+                } else if (extension === '.rs') {
+                    // For Rust files, tests are typically in the same file or in a tests module
+                    if (testFilesPath && testFilesPath.length > 0) {
+                        testFilePaths = [testFilesPath[0].fsPath];
+                    } else {
+                        // Check if the file already has a tests module
+                        const fileContent = fs.readFileSync(sourceFilePath, 'utf-8');
+                        const hasTestModule = fileContent.includes('#[cfg(test)]');
+
+                        if (!hasTestModule) {
+                            // Append test module to the source file
+                            const testModule = `
+
+                                #[cfg(test)]
+                                mod tests {
+                                    use super::*;
+
+                                    #[test]
+                                    fn test_dummy() {
+                                        assert!(true);
+                                    }
+                                }
+                                `;
+                            fs.appendFileSync(sourceFilePath, testModule);
+                            vscode.window.showInformationMessage('Added test module to the source file');
+                        }
+
+                        testFilePaths.push(sourceFilePath);
+                    }
+
+                    // Install required dependencies if not present
+                    if (!fs.existsSync(path.join(rootDir, 'Cargo.toml'))) {
+                        vscode.window.showErrorMessage('Cargo.toml not found. Please ensure this is a Rust project.');
+                        return;
+                    }
+
+                    // Check if grcov is installed
+                    exec('cargo install grcov', (error) => {
+                        if (error) {
+                            vscode.window.showErrorMessage('Failed to install grcov. Please install it manually: cargo install grcov');
+                        }
+                    });
+
+                    // Command to run tests with coverage using grcov
+                    command = `CARGO_INCREMENTAL=0 RUSTFLAGS="-Cinstrument-coverage" LLVM_PROFILE_FILE="cargo-test-%p-%m.profraw" cargo test && grcov . --binary-path ./target/debug/deps/ -s . -t cobertura --branch --ignore-not-existing --ignore "/*" -o coverage.xml`;
+                    coverageReportPath = "./coverage.xml";
+
                 } else {
                     vscode.window.showErrorMessage(`Unsupported file type: ${extension}`);
                     return;