feat: test (edit) command & add tests code lenses

lib/shared/src/commands/types.ts

@@ -4,13 +4,13 @@ export type DefaultCodyCommands = DefaultChatCommands | DefaultEditCommands
 // Default Cody Commands that runs as a Chat request
 export enum DefaultChatCommands {
     Explain = 'explain', // Explain code
-    Test = 'test', // Generate unit tests in Chat
+    Unit = 'unit', // Generate unit tests in Chat
     Smell = 'smell', // Generate code smell report in Chat
 }
 
 // Default Cody Commands that runs as an Inline Edit command
 export enum DefaultEditCommands {
-    Unit = 'unit', // Generate unit tests with inline edit
+    Test = 'test', // Generate unit tests with inline edit
     Doc = 'doc', // Generate documentation with inline edit
 }
 

vscode/package.json

@@ -287,16 +287,16 @@
       {
         "command": "cody.command.generate-tests",
         "category": "Cody Command",
-        "title": "Generate Unit Tests",
-        "icon": "$(package)",
+        "title": "Generate Unit Tests (Chat)",
+        "icon": "$(beaker)",
         "when": "cody.activated && editorTextFocus"
       },
       {
         "command": "cody.command.unit-tests",
         "category": "Cody Command",
-        "title": "Generate Unit Tests in File - Experimental",
+        "title": "Generate Unit Tests (Edit)",
         "icon": "$(package)",
-        "when": "cody.activated && config.cody.internal.unstable && editorTextFocus"
+        "when": "cody.activated && editorTextFocus"
       },
       {
         "command": "cody.command.document-code",
@@ -567,7 +567,7 @@
         },
         {
           "command": "cody.command.unit-tests",
-          "when": "cody.activated && config.cody.internal.unstable && editorIsOpen"
+          "when": "cody.activated && editorIsOpen"
         },
         {
           "command": "cody.command.document-code",
@@ -628,7 +628,7 @@
         },
         {
           "command": "cody.command.unit-tests",
-          "when": "cody.activated && config.cody.internal.unstable",
+          "when": "cody.activated",
           "group": "command"
         },
         {

vscode/src/commands/context/unit-test-case.ts

@@ -0,0 +1,46 @@
+import type { ContextFile } from '@sourcegraph/cody-shared'
+
+import { isTestFileForOriginal, isValidTestFile } from '../utils/test-commands'
+import { getWorkspaceFilesContext } from './workspace'
+import { getSearchPatternForTestFiles } from '../utils/search-pattern'
+import type { URI } from 'vscode-uri'
+import { getContextFileFromDirectory } from './directory'
+
+export async function getContextFilesForAddingUnitTestCases(testFile: URI): Promise<ContextFile[]> {
+    // Get the context from the current directory
+    // and then find the original file of the test file in the returned context
+    // If the original file is found, return it
+    // e.g. if the test file is src/foo/bar.spec.ts, look for src/foo/bar.ts
+    const directoryContext = await getContextFileFromDirectory()
+    const originalFileContext = directoryContext.find(f => isTestFileForOriginal(f.uri, testFile))
+    if (originalFileContext) {
+        return [originalFileContext]
+    }
+
+    // TODO (bee) improves context search
+    const contextFiles: ContextFile[] = []
+    // exclude any files in the path with e2e, integration, node_modules, or dist
+    const excludePattern = '**/*{e2e,integration,node_modules,dist}*/**'
+    // To search for files in the current directory only
+    const searchInCurrentDirectoryOnly = true
+    // The max number of files to search for in each workspace search
+    const max = 10
+
+    // Search for test files in the current directory first
+    const curerntDirPattern = getSearchPatternForTestFiles(testFile, searchInCurrentDirectoryOnly)
+    const currentDirContext = await getWorkspaceFilesContext(curerntDirPattern, excludePattern, max)
+
+    contextFiles.push(...currentDirContext)
+
+    // If no test files found in the current directory, search the entire workspace
+    if (!contextFiles.length) {
+        const wsTestPattern = getSearchPatternForTestFiles(testFile, !searchInCurrentDirectoryOnly)
+        // Will try to look for half the max number of files in the workspace for faster results
+        const codebaseFiles = await getWorkspaceFilesContext(wsTestPattern, excludePattern, max / 2)
+
+        contextFiles.push(...codebaseFiles)
+    }
+
+    // Return valid test files only
+    return contextFiles.filter(f => isValidTestFile(f.uri))
+}

vscode/src/commands/context/unit-test-command.ts → vscode/src/commands/context/unit-test-file.ts

@@ -24,14 +24,19 @@ export async function getContextFilesForUnitTestCommand(file: URI): Promise<Cont
     // To search for files in the current directory only
     const searchInCurrentDirectoryOnly = true
     // The max number of files to search for in each workspace search
-    const max = 5
+    const max = 10
 
+    // Search for test files in the current directory first
     const curerntDirPattern = getSearchPatternForTestFiles(file, searchInCurrentDirectoryOnly)
-    contextFiles.push(...(await getWorkspaceFilesContext(curerntDirPattern, excludePattern, max)))
+    const currentDirContext = await getWorkspaceFilesContext(curerntDirPattern, excludePattern, max)
 
+    contextFiles.push(...currentDirContext)
+
+    // If no test files found in the current directory, search the entire workspace
     if (!contextFiles.length) {
         const wsTestPattern = getSearchPatternForTestFiles(file, !searchInCurrentDirectoryOnly)
-        const codebaseFiles = await getWorkspaceFilesContext(wsTestPattern, excludePattern, max)
+        // Will try to look for half the max number of files in the workspace for faster results
+        const codebaseFiles = await getWorkspaceFilesContext(wsTestPattern, excludePattern, max / 2)
 
         contextFiles.push(...codebaseFiles)
     }

vscode/src/commands/execute/cody.json

@@ -20,25 +20,25 @@
             "type": "default"
         },
         "test": {
-            "description": "Generate unit tests",
-            "prompt": "Review the shared code context and configurations to identify the test framework and libraries in use. Then, generate a suite of multiple unit tests for the functions in <selected> using the detected test framework and libraries. Be sure to import the function being tested. Follow the same patterns as any shared context. Only add packages, imports, dependencies, and assertions if they are used in the shared code. Pay attention to the file path of each shared context to see if test for <selected> already exists. If one exists, focus on generating new unit tests for uncovered cases. If none are detected, import common unit test libraries for {languageName}. Focus on validating key functionality with simple and complete assertions. Only include mocks if one is detected in the shared code. Before writing the tests, identify which test libraries and frameworks to import, e.g. 'No new imports needed - using existing libs' or 'Importing test framework that matches shared context usage' or 'Importing the defined framework', etc. Then briefly summarize test coverage and any limitations. At the end, enclose the full completed code for the new unit tests, including all necessary imports, in a single markdown codeblock. No fragments or TODO. The new tests should validate expected functionality and cover edge cases for <selected> with all required imports, including importing the function being tested. Do not repeat existing tests.",
+            "description": "Generate unit tests (edit)",
+            "prompt": "Review the shared code context and configurations to identify the testing framework and libraries in use. Then, generate a suite of multiple unit tests for the functions in <selected> using the detected test framework and libraries. Be sure to import the function being tested. Follow the same patterns, testing conventions and testing library as shown in shared context. Only add packages, imports, dependencies, and assertions if they are used in the shared code. Pay attention to the file name of each shared context to see if test for <selected> already exists. If one exists, focus on generating new unit tests for uncovered cases. If none are detected, import common unit test libraries for {languageName}. Focus on validating key functionality with simple and complete assertions. Only include mocks if one is detected in the shared code. Before writing the tests, identify which testing libraries and frameworks to use and import. Then, enclose the file name for the unit test file between <FILEPATH7041> tags. At the end, enclose the full completed code for the new unit tests without any comments, fragments or TODO. The new tests should validate expected functionality and cover edge cases for <selected> with all required imports, including the function being tested. Do not repeat tests in shared context. Do not include any markdown formatting or triple backticks. Enclose only the complete runnable tests. If there is a shared context that has the same test file name, only include import if no conflict exists.",
             "context": {
                 "currentDir": true,
                 "currentFile": true,
                 "selection": true
             },
-            "type": "default"
+            "type": "default",
+            "mode": "file"
         },
         "unit": {
-            "description": "Experimental: Generate unit tests in test file directly",
-            "prompt": "Review the shared code context and configurations to identify the testing framework and libraries in use. Then, generate a suite of multiple unit tests for the functions in <selected> using the detected test framework and libraries. Be sure to import the function being tested. Follow the same patterns, testing conventions and testing library as shown in shared context. Only add packages, imports, dependencies, and assertions if they are used in the shared code. Pay attention to the file name of each shared context to see if test for <selected> already exists. If one exists, focus on generating new unit tests for uncovered cases. If none are detected, import common unit test libraries for {languageName}. Focus on validating key functionality with simple and complete assertions. Only include mocks if one is detected in the shared code. Before writing the tests, identify which testing libraries and frameworks to use and import. Then, enclose the file name for the unit test file between <FILEPATH7041> tags. At the end, enclose the full completed code for the new unit tests without any comments, fragments or TODO. The new tests should validate expected functionality and cover edge cases for <selected> with all required imports, including the function being tested. Do not repeat tests in shared context. Do not include any markdown formatting or triple backticks. Enclose only the complete runnable tests. If there is a shared context that has the same test file name, only include import if no conflict exists.",
+            "description": "Generate unit tests (chat)",
+            "prompt": "Review the shared code context and configurations to identify the test framework and libraries in use. Then, generate a suite of multiple unit tests for the functions in <selected> using the detected test framework and libraries. Be sure to import the function being tested. Follow the same patterns as any shared context. Only add packages, imports, dependencies, and assertions if they are used in the shared code. Pay attention to the file path of each shared context to see if test for <selected> already exists. If one exists, focus on generating new unit tests for uncovered cases. If none are detected, import common unit test libraries for {languageName}. Focus on validating key functionality with simple and complete assertions. Only include mocks if one is detected in the shared code. Before writing the tests, identify which test libraries and frameworks to import, e.g. 'No new imports needed - using existing libs' or 'Importing test framework that matches shared context usage' or 'Importing the defined framework', etc. Then briefly summarize test coverage and any limitations. At the end, enclose the full completed code for the new unit tests, including all necessary imports, in a single markdown codeblock. No fragments or TODO. The new tests should validate expected functionality and cover edge cases for <selected> with all required imports, including importing the function being tested. Do not repeat existing tests.",
             "context": {
                 "currentDir": true,
                 "currentFile": true,
                 "selection": true
             },
-            "type": "experimental",
-            "mode": "file"
+            "type": "default"
         },
         "smell": {
             "description": "Find code smells",

vscode/src/commands/execute/index.ts

@@ -5,18 +5,19 @@ import {
 } from '@sourcegraph/cody-shared/src/commands/types'
 import { executeSmellCommand } from './smell'
 import { executeExplainCommand } from './explain'
-import { executeTestCommand } from './test'
+import { executeUnitTestCommand } from './unit'
 import { executeDocCommand } from './doc'
 import type { CommandResult } from '../../main'
-import { executeUnitTestCommand } from './unit'
+import { executeTestCommand } from './test-file'
 
 export { commands as defaultCommands } from './cody.json'
 
 export { executeSmellCommand } from './smell'
 export { executeExplainCommand } from './explain'
-export { executeTestCommand } from './test'
-export { executeDocCommand } from './doc'
 export { executeUnitTestCommand } from './unit'
+export { executeDocCommand } from './doc'
+export { executeTestCommand } from './test-file'
+export { executeTestCaseCommand } from './test-case'
 
 export function isDefaultChatCommand(id: string): DefaultChatCommands | undefined {
     // Remove leading slash if any
@@ -51,10 +52,10 @@ export async function executeDefaultCommand(
             return executeExplainCommand({ additionalInstruction })
         case DefaultChatCommands.Smell:
             return executeSmellCommand({ additionalInstruction })
-        case DefaultChatCommands.Test:
-            return executeTestCommand({ additionalInstruction })
-        case DefaultEditCommands.Unit:
+        case DefaultChatCommands.Unit:
             return executeUnitTestCommand({ additionalInstruction })
+        case DefaultEditCommands.Test:
+            return executeTestCommand({ additionalInstruction })
         case DefaultEditCommands.Doc:
             return executeDocCommand({ additionalInstruction })
         default:

vscode/src/commands/execute/test-case.ts

@@ -0,0 +1,55 @@
+import { Range } from 'vscode'
+import { logError, type ContextFile } from '@sourcegraph/cody-shared'
+import { getEditor } from '../../editor/active-editor'
+import { type ExecuteEditArguments, executeEdit } from '../../edit/execute'
+import { DefaultEditCommands } from '@sourcegraph/cody-shared/src/commands/types'
+import type { EditCommandResult } from '../../main'
+import type { CodyCommandArgs } from '../types'
+import { getContextFilesForAddingUnitTestCases } from '../context/unit-test-case'
+
+/**
+ * NOTE: Used by CommandCodeLenses in test files with 'cody.command.unit-tests-cases'.
+ *
+ * It generates new test cases for the selected test suit.
+ * When invoked, the command will be executed as an inline-edit command.
+ */
+export async function executeTestCaseCommand(
+    args?: Partial<CodyCommandArgs>
+): Promise<EditCommandResult | undefined> {
+    const instruction =
+        'Review the shared code context to identify the testing framework and libraries in use. Then, create multiple new unit tests for the functions in <selected> following the same patterns, testing conventions and testing library as shown in shared context. Pay attention to the shared context to ensure your response code do not contain cases that have already been covered. Focus on generating new unit tests for uncovered cases. Response only with the full completed code with the new unit tests added at the end, without any comments, fragments or TODO. The new tests should validate expected functionality and cover edge cases for <selected>. Do not include any markdown formatting or triple backticks. The goal is to provide me with code that I can add to the end of existing test file. Enclose only the new tests WITHOUT ANY suit, import statements or packages in your response.'
+
+    const editor = getEditor()?.active
+    const document = editor?.document
+    // Current selection is required
+    if (!document || !editor.selection) {
+        return
+    }
+
+    const contextFiles: ContextFile[] = []
+    try {
+        const files = await getContextFilesForAddingUnitTestCases(document.uri)
+        contextFiles.push(...files)
+    } catch (error) {
+        logError('executeNewTestCommand', 'failed to fetch context', { verbose: error })
+    }
+
+    const startLine = editor.selection.start.line + 1
+    const endLine = Math.max(startLine, editor.selection.end.line)
+    const range = new Range(startLine, 0, endLine, 0)
+
+    return {
+        type: 'edit',
+        task: await executeEdit(
+            {
+                instruction,
+                document,
+                range,
+                intent: 'edit',
+                mode: 'insert',
+                userContextFiles: contextFiles,
+            } satisfies ExecuteEditArguments,
+            DefaultEditCommands.Doc
+        ),
+    }
+}