From ce63243ba06340c4460c7d186bfbe7d4dc9f2b5d Mon Sep 17 00:00:00 2001 From: akshat99812 <138353837+akshat99812@users.noreply.github.com> Date: Sat, 26 Oct 2024 12:04:48 +0530 Subject: [PATCH 1/2] rust-utg Signed-off-by: akshat99812 <138353837+akshat99812@users.noreply.github.com> --- package-lock.json | 20 +++++++++ package.json | 6 +-- scripts/utg.sh | 44 ++++++++++++++++++- src/Utg.ts | 82 ++++++++++++++++++++-------------- src/extension.ts | 110 +++++++++++++++++++++++++++++++++------------- 5 files changed, 193 insertions(+), 69 deletions(-) diff --git a/package-lock.json b/package-lock.json index d69fee3..769804c 100644 --- a/package-lock.json +++ b/package-lock.json @@ -19,6 +19,7 @@ "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" @@ -4669,6 +4670,25 @@ "resolved": "https://registry.npmjs.org/node-addon-api/-/node-addon-api-7.1.1.tgz", "integrity": "sha512-5m3bsyrjFWE1xf7nz7YXdN4udnVtXK6/Yfgn5qnahL6bCkf2yKt4k3nuTKAtT4r3IG8JNR2ncsIMdZuAzJjHQQ==" }, + "node_modules/tree-sitter-rust": { + "version": "0.23.0", + "resolved": "https://registry.npmjs.org/tree-sitter-rust/-/tree-sitter-rust-0.23.0.tgz", + "integrity": "sha512-vbb5ESloUtrRZMFA5EOgR56DMq6Ijb22PdG2zwsw8pdByYGLo0DxOmLp5KQce5rwNoZgEEp5/IiCN243TxQKIg==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.1.0", + "node-gyp-build": "^4.8.2" + }, + "peerDependencies": { + "tree-sitter": "^0.21.1" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, "node_modules/ts-api-utils": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/ts-api-utils/-/ts-api-utils-1.3.0.tgz", diff --git a/package.json b/package.json index fe3dd97..4548fd9 100644 --- a/package.json +++ b/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" diff --git a/scripts/utg.sh b/scripts/utg.sh index acd611c..8731695 100644 --- a/scripts/utg.sh +++ b/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 \ No newline at end of file diff --git a/src/Utg.ts b/src/Utg.ts index 5118f48..8425a71 100644 --- a/src/Utg.ts +++ b/src/Utg.ts @@ -215,45 +215,59 @@ async function Utg(context: vscode.ExtensionContext , additional_prompts?:string // **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; } - console.log("additional_prompts" , additional_prompts); - if(!additional_prompts){ - additional_prompts = ""; - } - - // Adjust the terminal command to include the test file path - terminal.sendText(`sh "${scriptPath}" "${sourceFilePath}" "${testFilePaths[0]}" "${coverageReportPath}" "${command}" "${additional_prompts}";`); - - const delay = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); - - // Add a 5-second delay before calling the API - await delay(5000); - - try { - if (token) { - apiResponse = await makeApiRequest(token) || 'no response'; - const response = JSON.parse(apiResponse); - await context.globalState.update('apiResponse', apiResponse); - if (response.usedCall === response.totalCall) { - await context.globalState.update('SubscriptionEnded', true); - } - } else { - console.log("token not found"); - } - } catch (apiError) { - vscode.window.showErrorMessage('Error during API request: ' + apiError); - } - - const disposable = vscode.window.onDidCloseTerminal(eventTerminal => { - if (eventTerminal === terminal) { - disposable.dispose(); - resolve(); - } - }); + // ... rest of the existing code ... } catch (error) { console.log(error); vscode.window.showErrorMessage('Error occurred Keploy utg: ' + error); diff --git a/src/extension.ts b/src/extension.ts index f22c4ed..7c73e99 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -1,5 +1,5 @@ import * as vscode from 'vscode'; -import * as path from 'path'; // Import path module +import * as path from 'path'; import { SidebarProvider } from './SidebarProvider'; import SignIn, { validateFirst, SignInWithOthers, ValidateSignInWithOthers } from './SignIn'; import oneClickInstall from './OneClickInstall'; @@ -13,45 +13,42 @@ import TreeSitterJavaScript from 'tree-sitter-javascript'; import TreeSitterPython from 'tree-sitter-python'; import TreeSitterJava from 'tree-sitter-java'; import TreeSitterGo from 'tree-sitter-go'; +import TreeSitterRust from 'tree-sitter-rust'; class KeployCodeLensProvider implements vscode.CodeLensProvider { onDidChangeCodeLenses?: vscode.Event | undefined; private treeCache: { [filePath: string]: TreeSitter.Tree } = {}; private invalidateCache(fileName: string) { - delete this.treeCache[fileName]; // Remove the cached tree + delete this.treeCache[fileName]; } constructor() { - // Listen for document changes and invalidates cache vscode.workspace.onDidChangeTextDocument(this.onDocumentChange.bind(this)); - - // Listen for document save events and invalidates cache vscode.workspace.onDidSaveTextDocument(this.onDocumentSave.bind(this)); } - // Triggered when the document is saved private onDocumentSave(document: vscode.TextDocument) { console.log('Document saved:', document.uri.fsPath); const fileName = document.uri.fsPath; - this.invalidateCache(fileName); // Invalidate the cache on save + this.invalidateCache(fileName); } private onDocumentChange(event: vscode.TextDocumentChangeEvent) { console.log('Document changed:', event.document.uri.fsPath); const fileName = event.document.uri.fsPath; - this.invalidateCache(fileName); // Invalidate the cache on change + this.invalidateCache(fileName); } private getTreeFromCache(fileName: string, text: string): TreeSitter.Tree { if (this.treeCache[fileName]) { console.log(`Cache hit for: ${fileName}`); - this.logCacheSize(); // Log the cache size - return this.treeCache[fileName]; // Return cached tree if available + this.logCacheSize(); + return this.treeCache[fileName]; } const parser = new TreeSitter(); - // Set the language based on file extension + // Add Rust support to language selection if (fileName.endsWith('.js') || fileName.endsWith('.ts')) { parser.setLanguage(TreeSitterJavaScript); } else if (fileName.endsWith('.py')) { @@ -60,18 +57,21 @@ class KeployCodeLensProvider implements vscode.CodeLensProvider { parser.setLanguage(TreeSitterJava); } else if (fileName.endsWith('.go')) { parser.setLanguage(TreeSitterGo); + } else if (fileName.endsWith('.rs')) { + parser.setLanguage(TreeSitterRust); } else { console.log('Unsupported file type:', fileName); throw new Error("Unsupported file type"); } + const options: TreeSitter.Options = { bufferSize: 1024 * 1024, }; - const tree = parser.parse(text, undefined, options); // Parse the document text - this.treeCache[fileName] = tree; // Cache the parsed tree + const tree = parser.parse(text, undefined, options); + this.treeCache[fileName] = tree; console.log(`Cache miss for: ${fileName}`); - this.logCacheSize(); // Log the cache size after adding a new entry + this.logCacheSize(); return tree; } @@ -118,25 +118,27 @@ class KeployCodeLensProvider implements vscode.CodeLensProvider { ): vscode.CodeLens[] | Thenable { const fileName = document.uri.fsPath; - // Skip test files + // Add Rust test file pattern 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 + fileName.endsWith('Tests.java') || + fileName.endsWith('Test.java') || + fileName.includes('/Test') || + fileName.includes('/test/') || fileName.endsWith('_test.go') || - fileName.includes('test_') + fileName.includes('test_') || + fileName.endsWith('_test.rs') ) { return []; } + const baseName = path.basename(fileName); const fileExtension = path.extname(fileName); - const text = document.getText(); const codeLenses: vscode.CodeLens[] = []; + try { const tree = this.getTreeFromCache(fileName, text); const cursor = tree.walk(); @@ -162,6 +164,24 @@ class KeployCodeLensProvider implements vscode.CodeLensProvider { command: 'keploy.showSidebar', arguments: [document.uri.fsPath, functionName, fileExtension] })); + console.log('🐰 Found arrow function:', node.firstChild?.text); + }else if (fileName.endsWith('.rs') && node.type === 'function_item') { + const line = document.positionAt(node.startIndex).line; + const range = new vscode.Range(line, 0, line, 0); + const functionName = node.childForFieldName('name')?.text || ''; + + codeLenses.push(new vscode.CodeLens(range, { + title: '🐰 Generate unit tests', + command: 'keploy.utg', + arguments: [document.uri.fsPath, functionName, fileExtension] + })); + + codeLenses.push(new vscode.CodeLens(range, { + title: '🐰 Additional Prompts', + command: 'keploy.showSidebar', + arguments: [document.uri.fsPath, functionName, fileExtension] + })); + console.log('🐰 Found arrow function:', node.firstChild?.text); } else if (fileName.endsWith('.js') || fileName.endsWith('.ts')) { if (node.type === 'arrow_function') { @@ -258,25 +278,23 @@ class KeployCodeLensProvider implements vscode.CodeLensProvider { } async function findTestCasesForFunction(functionName: string, fileExtension: string): Promise { console.log(`🐰 Searching for test cases for function: ${functionName}`); - // Exclude certain directories from the search - const excludePattern = '**/{node_modules,venv,__pycache__}/**'; // Exclude common directories - const maxResults = 100; // Limit the number of files returned - console.log('🐰 Searching for test files...'); + const excludePattern = '**/{node_modules,venv,__pycache__,target}/**'; + const maxResults = 100; + let testFilePattern = ''; if (fileExtension === '.js' || fileExtension === '.ts') { testFilePattern = '**/*.{test,spec}.{js,ts}'; } else if (fileExtension === '.py') { testFilePattern = '**/{test_*,*_test}.py'; + } else if (fileExtension === '.rs') { + testFilePattern = '**/*_test.rs'; } else { - // Unsupported language for now return; } - // Use await to get the list of test files const testFiles = await vscode.workspace.findFiles(testFilePattern, excludePattern, maxResults); console.log('🐰 Found test files:', testFiles); - // Array to collect found test files const foundTestFiles: vscode.Uri[] = []; for (const fileUri of testFiles) { @@ -288,10 +306,14 @@ async function findTestCasesForFunction(functionName: string, fileExtension: str try { const testFileExtension = fileUri.fsPath.split('.').pop(); const parser = new TreeSitter(); + + // Add Rust parser support if (testFileExtension === 'js' || testFileExtension === 'ts') { parser.setLanguage(TreeSitterJavaScript); } else if (testFileExtension === 'py') { parser.setLanguage(TreeSitterPython); + } else if (testFileExtension === 'rs') { + parser.setLanguage(TreeSitterRust); } else { console.log('Unsupported test file language:', testFileExtension); continue; @@ -338,6 +360,21 @@ async function findTestCasesForFunction(functionName: string, fileExtension: str found = true; } + }else if (testFileExtension === 'rs') { + // Check for mod test blocks + if (node.type === 'mod_item' && node.childForFieldName('name')?.text === 'tests') { + found = true; + } + + // Check for test functions + if (node.type === 'function_item' && + node.previousNamedSibling?.type === 'attribute' && + node.previousNamedSibling.text.includes('#[test]')) { + if (node.text.includes(functionName)) { + found = true; + console.log(`🐰 Test function for ${functionName} found in file: ${fileUri.fsPath}`); + } + } } else if (testFileExtension === 'py') { // Check for function calls if (node.type === 'call') { @@ -407,6 +444,7 @@ async function getAllFunctionsInFile( const text = document.getText(); const parser = new TreeSitter(); + // Add Rust parser support if (fileExtension === '.js' || fileExtension === '.ts') { parser.setLanguage(TreeSitterJavaScript); } else if (fileExtension === '.py') { @@ -415,16 +453,18 @@ async function getAllFunctionsInFile( parser.setLanguage(TreeSitterJava); } else if (fileExtension === '.go') { parser.setLanguage(TreeSitterGo); + } else if (fileExtension === '.rs') { + parser.setLanguage(TreeSitterRust); } else { 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, undefined, options); const cursor = tree.walk(); const functionNames: string[] = []; @@ -440,6 +480,12 @@ async function getAllFunctionsInFile( functionNames.push(functionName); console.log(`🐰 Found function: ${functionName}`); } + }else if (fileExtension === '.rs' && node.type === 'function_item') { + const functionName = node.childForFieldName('name')?.text; + if (functionName) { + functionNames.push(functionName); + console.log(`🐰 Found function: ${functionName}`); + } } else if (fileExtension === '.py' && node.type === 'function_definition') { const functionName = node.childForFieldName('name')?.text; if (functionName) { @@ -509,6 +555,10 @@ export function activate(context: vscode.ExtensionContext) { "Keploy-Sidebar", sidebarProvider ), + vscode.languages.registerCodeLensProvider( + { language: 'rust', scheme: 'file' }, + new KeployCodeLensProvider() + ), vscode.languages.registerCodeLensProvider( { language: 'javascript', scheme: 'file' }, new KeployCodeLensProvider() From 94829f0ec97e50784ac891839ac5de8cc2cd1403 Mon Sep 17 00:00:00 2001 From: akshat99812 <138353837+akshat99812@users.noreply.github.com> Date: Sat, 26 Oct 2024 12:37:06 +0530 Subject: [PATCH 2/2] utf-fix Signed-off-by: akshat99812 <138353837+akshat99812@users.noreply.github.com> --- src/Utg.ts | 35 ++++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/src/Utg.ts b/src/Utg.ts index 8425a71..1fb871b 100644 --- a/src/Utg.ts +++ b/src/Utg.ts @@ -267,7 +267,40 @@ async function Utg(context: vscode.ExtensionContext , additional_prompts?:string return; } - // ... rest of the existing code ... + console.log("additional_prompts" , additional_prompts); + if(!additional_prompts){ + additional_prompts = ""; + } + + // Adjust the terminal command to include the test file path + terminal.sendText(`sh "${scriptPath}" "${sourceFilePath}" "${testFilePaths[0]}" "${coverageReportPath}" "${command}" "${additional_prompts}";`); + + const delay = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); + + // Add a 5-second delay before calling the API + await delay(5000); + + try { + if (token) { + apiResponse = await makeApiRequest(token) || 'no response'; + const response = JSON.parse(apiResponse); + await context.globalState.update('apiResponse', apiResponse); + if (response.usedCall === response.totalCall) { + await context.globalState.update('SubscriptionEnded', true); + } + } else { + console.log("token not found"); + } + } catch (apiError) { + vscode.window.showErrorMessage('Error during API request: ' + apiError); + } + + const disposable = vscode.window.onDidCloseTerminal(eventTerminal => { + if (eventTerminal === terminal) { + disposable.dispose(); + resolve(); + } + }); } catch (error) { console.log(error); vscode.window.showErrorMessage('Error occurred Keploy utg: ' + error);