Revert "Agent: remove deprecated unit test command (#5718)"

agent/src/agent.ts

@@ -1090,6 +1090,18 @@ export class Agent extends MessageHandler implements ExtensionClient {
             )
         })
 
+        this.registerAuthenticatedRequest('commands/test', () => {
+            return this.createChatPanel(
+                vscode.commands.executeCommand('cody.command.generate-tests', commandArgs)
+            )
+        })
+
+        this.registerAuthenticatedRequest('editCommands/test', () => {
+            return this.createEditTask(
+                vscode.commands.executeCommand<CommandResult | undefined>('cody.command.unit-tests')
+            )
+        })
+
         this.registerAuthenticatedRequest('editTask/accept', async ({ id }) => {
             this.fixups?.accept(id)
             return null

agent/src/index.test.ts

@@ -9,6 +9,7 @@
     DOTCOM_URL,
     ModelUsage,
     type SerializedChatTranscript,
+    isWindows,
 } from '@sourcegraph/cody-shared'
 
 import * as uuid from 'uuid'
@@ -837,6 +838,29 @@
             )
         }, 30_000)
 
+        // This test seems extra sensitive on Node v16 for some reason.
+        it.skipIf(isWindows())(
+            'commands/test',
+            async () => {
+                await client.openFile(animalUri)
+                await setChatModel()
+                const id = await client.request('commands/test', null)
+                const lastMessage = await client.firstNonEmptyTranscript(id)
+                expect(trimEndOfLine(lastMessage.messages.at(-1)?.text ?? '')).toMatchSnapshot()
+                // telemetry assertion, to validate the expected events fired during the test run
+                // Do not remove this assertion, and instead update the expectedEvents list above
+                expect(await exportedTelemetryEvents(client)).toEqual(
+                    expect.arrayContaining([
+                        'cody.command.test:executed',
+                        'cody.chat-question:submitted',
+                        'cody.chat-question:executed',
+                        'cody.chatResponse:hasCode',
+                    ])
+                )
+            },
+            30_000
+        )
+
         it('commands/smell', async () => {
             await client.openFile(animalUri)
             await setChatModel()

lib/shared/src/commands/types.ts

@@ -5,6 +5,7 @@ export type DefaultCodyCommands = DefaultChatCommands | DefaultEditCommands
 export enum DefaultChatCommands {
     Doc = 'doc', // Generate documentation in Chat
     Explain = 'explain', // Explain code
+    Unit = 'unit', // Generate unit tests in Chat
     Smell = 'smell', // Generate code smell report in Chat
     Custom = 'custom-chat', // Run custom command in Chat
 }

vscode/src/commands/CommandsController.ts

@@ -127,6 +127,8 @@ function convertDefaultCommandsToPromptString(input: DefaultCodyCommands | Promp
             return ps`explain`
         case DefaultChatCommands.Smell:
             return ps`smell`
+        case DefaultChatCommands.Unit:
+            return ps`unit`
         case DefaultEditCommands.Test:
             return ps`test`
         case DefaultChatCommands.Doc:

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

@@ -0,0 +1,45 @@
+import { type ContextItem, wrapInActiveSpan } from '@sourcegraph/cody-shared'
+
+import type * as vscode from 'vscode'
+
+import { getSearchPatternForTestFiles } from '../utils/search-pattern'
+import { isValidTestFile } from '../utils/test-commands'
+import { getContextFileFromDirectory } from './directory'
+import { getWorkspaceFilesContext } from './workspace'
+
+/**
+ * Gets context files related to the given test file.
+ *
+ * Searches for test files in the current directory first.
+ * If none found, searches the entire workspace for test files.
+ *
+ * Returns only valid test files up to the max limit.
+ *
+ * NOTE: This is used by the current unit test commands to get context files.
+ * NOTE: Will be replaced by the new unit test commands once it's ready.
+ */
+export async function getContextFilesForTestCommand(file: vscode.Uri): Promise<ContextItem[]> {
+    return wrapInActiveSpan('commands.context.testChat', async span => {
+        const contextFiles: ContextItem[] = []
+
+        // 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 = 5
+
+        // Get context from test files in current directory
+        contextFiles.push(...(await getContextFileFromDirectory()))
+
+        if (!contextFiles.length) {
+            const wsTestPattern = getSearchPatternForTestFiles(file, !searchInCurrentDirectoryOnly)
+            const codebaseFiles = await getWorkspaceFilesContext(wsTestPattern, excludePattern, max)
+
+            contextFiles.push(...codebaseFiles)
+        }
+
+        // Return valid test files only
+        return contextFiles.filter(f => isValidTestFile(f.uri))
+    })
+}

vscode/src/commands/execute/index.ts

@@ -11,12 +11,14 @@ import { executeEdit } from '../../edit/execute'
 import { executeDocCommand } from './doc'
 import { executeExplainCommand } from './explain'
 import { executeSmellCommand } from './smell'
+import { executeTestChatCommand } from './test-chat'
 import { executeTestEditCommand } from './test-edit'
 
 export { commands as defaultCommands } from './cody.json'
 
 export { executeSmellCommand } from './smell'
 export { executeExplainCommand } from './explain'
+export { executeTestChatCommand } from './test-chat'
 export { executeDocCommand } from './doc'
 export { executeTestEditCommand } from './test-edit'
 export { executeTestCaseEditCommand } from './test-case'
@@ -55,6 +57,8 @@ export async function executeDefaultCommand(
             return executeExplainCommand({ additionalInstruction })
         case DefaultChatCommands.Smell:
             return executeSmellCommand({ additionalInstruction })
+        case DefaultChatCommands.Unit:
+            return executeTestChatCommand({ additionalInstruction })
         case DefaultEditCommands.Test:
             return executeTestEditCommand({ additionalInstruction })
         case DefaultEditCommands.Doc:

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

@@ -0,0 +1,111 @@
+import { type ContextItem, DefaultChatCommands, logDebug, logError, ps } from '@sourcegraph/cody-shared'
+import { wrapInActiveSpan } from '@sourcegraph/cody-shared'
+import { telemetryRecorder } from '@sourcegraph/cody-shared'
+import type { ChatCommandResult } from '../../CommandResult'
+import { getEditor } from '../../editor/active-editor'
+import { getContextFileFromCursor } from '../context/selection'
+import { getContextFilesForTestCommand } from '../context/unit-test-chat'
+import type { CodyCommandArgs } from '../types'
+import { type ExecuteChatArguments, executeChat } from './ask'
+
+import type { Span } from '@opentelemetry/api'
+import { isUriIgnoredByContextFilterWithNotification } from '../../cody-ignore/context-filter'
+import { selectedCodePromptWithExtraFiles } from './index'
+
+/**
+ * Generates the prompt and context files with arguments for the '/test' command in Chat.
+ *
+ * Context: Test files, current selection, and current file
+ */
+async function unitTestCommand(
+    span: Span,
+    args?: Partial<CodyCommandArgs>
+): Promise<ExecuteChatArguments> {
+    let prompt = ps`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.`
+
+    if (args?.additionalInstruction) {
+        prompt = ps`${prompt} ${args.additionalInstruction}`
+    }
+
+    const editor = getEditor()?.active
+    const document = editor?.document
+    const contextItems: ContextItem[] = []
+
+    if (document) {
+        try {
+            const cursorContext = await getContextFileFromCursor()
+            if (cursorContext === null) {
+                throw new Error(
+                    'Selection content is empty. Please select some code to generate tests for.'
+                )
+            }
+
+            const sharedContext = await getContextFilesForTestCommand(document.uri)
+
+            prompt = prompt.replaceAll('<selected>', selectedCodePromptWithExtraFiles(cursorContext, []))
+
+            if (sharedContext.length > 0) {
+                prompt = prompt.replaceAll(
+                    'the shared code',
+                    selectedCodePromptWithExtraFiles(sharedContext[0], sharedContext.slice(1))
+                )
+            }
+
+            contextItems.push(cursorContext)
+            contextItems.push(...sharedContext)
+        } catch (error) {
+            logError('testCommand', 'failed to fetch context', { verbose: error })
+        }
+    }
+
+    return {
+        text: prompt,
+        contextItems,
+        source: args?.source,
+        submitType: 'user-newchat',
+        command: DefaultChatCommands.Unit,
+    }
+}
+
+/**
+ * Executes the /test command for generating unit tests in Chat for selected code.
+ *
+ * NOTE: Currently used by agent until inline test command is added to agent.
+ */
+export async function executeTestChatCommand(
+    args?: Partial<CodyCommandArgs>
+): Promise<ChatCommandResult | undefined> {
+    return wrapInActiveSpan('command.test-chat', async span => {
+        span.setAttribute('sampled', true)
+
+        const editor = getEditor()
+        if (
+            editor.active &&
+            (await isUriIgnoredByContextFilterWithNotification(editor.active.document.uri, 'test'))
+        ) {
+            return
+        }
+
+        logDebug('executeTestEditCommand', 'executing', { args })
+        telemetryRecorder.recordEvent('cody.command.test', 'executed', {
+            metadata: {
+                useCodebaseContex: 0,
+            },
+            interactionID: args?.requestID,
+            privateMetadata: {
+                requestID: args?.requestID,
+                source: args?.source,
+                traceId: span.spanContext().traceId,
+            },
+            billingMetadata: {
+                product: 'cody',
+                category: 'core',
+            },
+        })
+
+        return {
+            type: 'chat',
+            session: await executeChat(await unitTestCommand(span, args)),
+        }
+    })
+}

vscode/src/jsonrpc/agent-protocol.ts

@@ -103,6 +103,7 @@ export type ClientRequests = {
     // of these commands is the same as `chat/new`, an ID to reference to the
     // webview panel where the reply from this command appears.
     'commands/explain': [null, string] // TODO: rename to chatCommands/{explain,test,smell}
+    'commands/test': [null, string]
     'commands/smell': [null, string]
 
     // Trigger custom commands that could be a chat-based command or an edit command.

vscode/src/main.ts

@@ -59,6 +59,7 @@ import {
     executeExplainOutput,
     executeSmellCommand,
     executeTestCaseEditCommand,
+    executeTestChatCommand,
     executeTestEditCommand,
 } from './commands/execute'
 import { executeAutoEditCommand } from './commands/execute/auto-edit'
@@ -439,6 +440,9 @@ async function registerCodyCommands(
                                   vscode.commands.registerCommand('cody.command.document-code', a =>
                                       executeDocChatCommand(a)
                                   ),
+                                  vscode.commands.registerCommand('cody.command.unit-tests', a =>
+                                      executeTestChatCommand(a)
+                                  ),
                               ]
                             : [
                                   // Otherwise register old-style commands.
@@ -454,6 +458,9 @@ async function registerCodyCommands(
                                   vscode.commands.registerCommand('cody.command.document-code', a =>
                                       executeDocCommand(a)
                                   ),
+                                  vscode.commands.registerCommand('cody.command.generate-tests', a =>
+                                      executeTestChatCommand(a)
+                                  ),
                                   vscode.commands.registerCommand('cody.command.unit-tests', a =>
                                       executeTestEditCommand(a)
                                   ),