Skip to content

Commit

Permalink
feat(bench): unit test framework for cody-bench (#4710)
Browse files Browse the repository at this point in the history 
Co-authored-by: Julie Tibshirani <[email protected]>
  • Loading branch information
@jamesmcnamara @jtibshirani
jamesmcnamara and jtibshirani authored Jul 1, 2024
1 parent bec6457 commit 6ebe1f1
Show file tree
Hide file tree
Showing 6 changed files with 326 additions and 40 deletions.
8 changes: 6 additions & 2 deletions agent/src/AgentTextDocument.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,12 @@ export class AgentTextDocument implements vscode.TextDocument {
public version = 0
public readonly isDirty: boolean = false
public readonly isClosed: boolean = false
public static from(uri: vscode.Uri, content: string): AgentTextDocument {
return new AgentTextDocument(ProtocolTextDocumentWithUri.from(uri, { content }))
public static from(
uri: vscode.Uri,
content: string,
document?: Partial<protocol.ProtocolTextDocument>
): AgentTextDocument {
return new AgentTextDocument(ProtocolTextDocumentWithUri.from(uri, { ...document, content }))
}

public save(): Thenable<boolean> {
Expand Down
117 changes: 94 additions & 23 deletions agent/src/TestClient.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,12 @@ import type {
EditTask,
ExtensionConfiguration,
NetworkRequest,
Position,
ProgressReportParams,
ProgressStartParams,
ProtocolCodeLens,
ProtocolTextDocument,
Range,
RenameFileOperation,
ServerInfo,
ShowWindowMessageParams,
Expand Down Expand Up @@ -270,11 +272,8 @@ export class TestClient extends MessageHandler {
const createdFiles: CreateFileOperation[] = []
for (const operation of params.operations) {
if (operation.type === 'edit-file') {
const { success, protocolDocument } = this.editDocument(operation)
result ||= success
if (protocolDocument) {
this.notify('textDocument/didChange', protocolDocument.underlying)
}
const protocolDocument = await this.editDocument(operation)
this.notify('textDocument/didChange', protocolDocument.underlying)
} else if (operation.type === 'create-file') {
const fileExists = await doesFileExist(vscode.Uri.parse(operation.uri))
if (operation.options?.ignoreIfExists && fileExists) {
Expand Down Expand Up @@ -343,13 +342,11 @@ export class TestClient extends MessageHandler {
this.notify('textDocument/didOpen', params)
return Promise.resolve(true)
})
this.registerRequest('textDocument/edit', params => {
this.registerRequest('textDocument/edit', async params => {
this.textDocumentEditParams.push(params)
const { success, protocolDocument } = this.editDocument(params)
if (protocolDocument) {
this.notify('textDocument/didChange', protocolDocument.underlying)
}
return Promise.resolve(success)
const protocolDocument = await this.editDocument(params)
this.notify('textDocument/didChange', protocolDocument.underlying)
return Promise.resolve(true)
})
this.registerRequest('textDocument/show', () => {
return Promise.resolve(true)
Expand All @@ -374,15 +371,8 @@ export class TestClient extends MessageHandler {
})
}

private editDocument(params: TextDocumentEditParams): {
success: boolean
protocolDocument?: ProtocolTextDocumentWithUri
} {
const document = this.workspace.getDocument(vscode.Uri.parse(params.uri))
if (!document) {
logError('textDocument/edit: document not found', params.uri)
return { success: false }
}
private async editDocument(params: TextDocumentEditParams): Promise<ProtocolTextDocumentWithUri> {
const document = await this.workspace.openTextDocument(vscode.Uri.parse(params.uri))
const patches = params.edits.map<[number, number, string]>(edit => {
switch (edit.type) {
case 'delete':
Expand All @@ -406,7 +396,7 @@ export class TestClient extends MessageHandler {
content: updatedContent,
})
this.workspace.loadDocument(protocolDocument)
return { success: true, protocolDocument }
return protocolDocument
}
private logMessage(params: DebugMessage): void {
// Uncomment below to see `logDebug` messages.
Expand Down Expand Up @@ -445,7 +435,7 @@ export class TestClient extends MessageHandler {
content = content.replace('/* CURSOR */', '')
}

const document = AgentTextDocument.from(uri, content)
const document = AgentTextDocument.from(uri, content, params)
const start =
cursor >= 0
? document.positionAt(cursor)
Expand All @@ -457,7 +447,7 @@ export class TestClient extends MessageHandler {
const protocolDocument: ProtocolTextDocument = {
uri: uri.toString(),
content,
selection: start && end ? { start, end } : undefined,
selection: params?.selection ?? (start && end ? { start, end } : undefined),
}
const clientDocument = this.workspace.loadDocument(
ProtocolTextDocumentWithUri.fromDocument(protocolDocument)
Expand Down Expand Up @@ -494,6 +484,70 @@ export class TestClient extends MessageHandler {
return this.workspace.getDocument(uri)?.content ?? ''
}

public async generateUnitTestFor(uri: vscode.Uri, line: number): Promise<TestInfo | undefined> {
await this.openFile(uri, {
removeCursor: false,
selection: {
start: { line, character: 0 },
end: { line, character: 0 },
},
})

const task = await this.request('editCommands/test', null)
await this.taskHasReachedAppliedPhase(task)
const lenses = this.codeLenses.get(uri.toString()) ?? []
if (lenses.length > 0) {
throw new Error(
`Code lenses are not supported in this mode ${JSON.stringify(lenses, null, 2)}`
)
}

await this.request('editTask/accept', { id: task.id })
return this.getTestEdit()
}

private async getTestEdit(): Promise<TestInfo | undefined> {
// first check if a new text file was created in the workspace
if (this.textDocumentEditParams.length === 1) {
const [editParams] = this.textDocumentEditParams
for (const edit of editParams.edits) {
if (edit.type === 'replace') {
return {
uri: vscode.Uri.parse(editParams.uri).with({ scheme: 'file' }),
value: edit.value,
fullFile: edit.value,
isUpdate: false,
}
}
}
}
// Otherwise it is an update to test file so it should
for (const param of this.workspaceEditParams) {
for (const operation of param.operations) {
if (operation.type === 'edit-file') {
for (const edit of operation.edits) {
// looks for a replace with an appropriate range
if (
edit.type === 'replace' &&
!isPositionEqual(edit.range.start, edit.range.end)
) {
const uri = vscode.Uri.parse(operation.uri).with({ scheme: 'file' })
await this.openFile(uri)
return {
uri,
value: edit.value,
fullFile: this.documentText(uri),
isUpdate: true,
}
}
}
}
}
}

return undefined
}

public async autocompleteText(params?: Partial<AutocompleteParams>): Promise<string[]> {
const result = await this.autocomplete(params)
return result.items.map(item => item.insertText)
Expand Down Expand Up @@ -952,4 +1006,21 @@ interface TextDocumentEventParams {
text?: string
selectionName?: string
removeCursor?: boolean
selection?: Range | undefined | null
}

function isPositionEqual(a: Position, b: Position): boolean {
return a.line === b.line && a.character === b.character
}

interface TestInfo {
// Test files "on disk" URI (not actually written about but uses "file" protocol)
uri: vscode.Uri

// Test content
value: string
fullFile: string

// Was this an update to an exisiting test file or a new file
isUpdate: boolean
}
26 changes: 25 additions & 1 deletion agent/src/cli/cody-bench/EvaluationDocument.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,8 @@ export class EvaluationDocument {
out.push(' FIX')
} else if (this.options.fixture.strategy === BenchStrategy.Chat) {
out.push(' CHAT')
} else if (this.options.fixture.strategy === BenchStrategy.UnitTest) {
out.push(' UNIT TEST')
} else {
throw new Error(`unknown strategy ${this.options.fixture.strategy}`)
}
Expand Down Expand Up @@ -272,9 +274,22 @@ interface EvaluationItem {
info?: CompletionItemInfo
event?: CompletionBookkeepingEvent
eventJSON?: string
testName?: string
testExpectedFilename?: string
testFilename?: string
testInputFilename?: string
testLanguage?: string
testGenerated?: string
testUsedExpectedTestFramework?: boolean
testUsedCorrectAppendOperation?: boolean
testDiagnostics?: string
}

export const headerItems: ObjectHeaderItem[] = [
interface EvaluationItemHeader extends ObjectHeaderItem {
id: keyof EvaluationItem
}

export const headerItems: EvaluationItemHeader[] = [
{ id: 'languageid', title: 'LANGUAGEID' },
{ id: 'workspace', title: 'WORKSPACE' },
{ id: 'fixture', title: 'FIXTURE' },
Expand Down Expand Up @@ -314,6 +329,15 @@ export const headerItems: ObjectHeaderItem[] = [
{ id: 'contextBfgSuggestedCount', title: 'CONTEXT_BFG_SUGGESTED_COUNT' },
{ id: 'contextBfgDurationMs', title: 'CONTEXT_BFG_DURATION_MS' },
{ id: 'eventJSON', title: 'EVENT' },
{ id: 'testFilename', title: 'TEST_FILENAME' },
{ id: 'testExpectedFilename', title: 'TEST_EXPECTED_FILENAME' },
{ id: 'testGenerated', title: 'TEST_GENERATED' },
{ id: 'testUsedExpectedTestFramework', title: 'TEST_USED_EXPECTED_TEST_FRAMEWORK' },
{ id: 'testUsedCorrectAppendOperation', title: 'TEST_USED_CORRECT_APPEND_OPERATION' },
{ id: 'testInputFilename', title: 'TEST_INPUT_FILENAME' },
{ id: 'testLanguage', title: 'TEST_LANGUAGE' },
{ id: 'testName', title: 'TEST_NAME' },
{ id: 'testDiagnostics', title: 'TEST_DIAGNOSTICS' },
]

function commentSyntaxForLanguage(languageid: string): string {
Expand Down
69 changes: 55 additions & 14 deletions agent/src/cli/cody-bench/cody-bench.ts
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import * as fspromises from 'node:fs/promises'
import * as path from 'node:path'
import glob from 'glob'

import * as commander from 'commander'
import * as vscode from 'vscode'
Expand All @@ -8,16 +9,19 @@ import { newAgentClient } from '../../agent'

import { ModelsService, getDotComDefaultModels, graphqlClient } from '@sourcegraph/cody-shared'
import { startPollyRecording } from '../../../../vscode/src/testutils/polly'
import { dotcomCredentials } from '../../../../vscode/src/testutils/testing-credentials'
import { allClientCapabilitiesEnabled } from '../../allClientCapabilitiesEnabled'
import { arrayOption, booleanOption, intOption } from './cli-parsers'
import { matchesGlobPatterns } from './matchesGlobPatterns'
import { evaluateAutocompleteStrategy } from './strategy-autocomplete'
import { evaluateChatStrategy } from './strategy-chat'
import { evaluateFixStrategy } from './strategy-fix'
import { evaluateGitLogStrategy } from './strategy-git-log'
import { evaluateUnitTestStrategy } from './strategy-unit-test'

export interface CodyBenchOptions {
workspace: string
absolutePath?: string
worktree?: string
treeSitterGrammars: string
queriesDirectory: string
Expand Down Expand Up @@ -64,9 +68,10 @@ interface EvaluationConfig extends Partial<CodyBenchOptions> {

export enum BenchStrategy {
Autocomplete = 'autocomplete',
GitLog = 'git-log',
Fix = 'fix',
Chat = 'chat',
Fix = 'fix',
GitLog = 'git-log',
UnitTest = 'unit-test',
}

interface EvaluationFixture {
Expand All @@ -83,14 +88,14 @@ async function loadEvaluationConfig(options: CodyBenchOptions): Promise<CodyBenc
const configBuffer = await fspromises.readFile(options.evaluationConfig)
const config = JSON.parse(configBuffer.toString()) as EvaluationConfig
const result: CodyBenchOptions[] = []
for (const test of config?.workspaces ?? []) {
const rootDir = path.dirname(options.evaluationConfig)
for (const test of expandWorkspaces(config.workspaces, rootDir) ?? []) {
if (!test.workspace) {
console.error(
`skipping test, missing required property 'workspace': ${JSON.stringify(test)}`
)
continue
}
const rootDir = path.dirname(options.evaluationConfig)
const workspace = path.normalize(path.join(rootDir, test.workspace))
const fixtures: EvaluationFixture[] = config.fixtures ?? [
{ name: 'default', strategy: BenchStrategy.Autocomplete },
Expand Down Expand Up @@ -202,7 +207,9 @@ export const codyBenchCommand = new commander.Command('cody-bench')
new commander.Option(
'--src-endpoint <url>',
'The Sourcegraph URL endpoint to use for authentication'
).env('SRC_ENDPOINT')
)
.env('SRC_ENDPOINT')
.default('https://sourcegraph.com')
)
.option(
'--include-workspace <glob>',
Expand Down Expand Up @@ -278,6 +285,19 @@ export const codyBenchCommand = new commander.Command('cody-bench')
true
)
.action(async (options: CodyBenchOptions) => {
if (!options.srcAccessToken) {
const { token } = dotcomCredentials()
if (!token) {
console.error('environment variable SRC_ACCESS_TOKEN must be non-empty')
process.exit(1)
}
options.srcAccessToken = token
}
if (!options.srcEndpoint) {
console.error('environment variable SRC_ENDPOINT must be non-empty')
process.exit(1)
}

const testOptions = await loadEvaluationConfig(options)
const workspacesToRun = testOptions.filter(
testOptions =>
Expand Down Expand Up @@ -322,15 +342,6 @@ export const codyBenchCommand = new commander.Command('cody-bench')
async function evaluateWorkspace(options: CodyBenchOptions, recordingDirectory: string): Promise<void> {
console.log(`starting evaluation: fixture=${options.fixture.name} workspace=${options.workspace}`)

if (!options.srcAccessToken) {
console.error('environment variable SRC_ACCESS_TOKEN must be non-empty')
process.exit(1)
}
if (!options.srcEndpoint) {
console.error('environment variable SRC_ENDPOINT must be non-empty')
process.exit(1)
}

const workspaceRootUri = vscode.Uri.from({ scheme: 'file', path: options.workspace })

const baseGlobalState: Record<string, any> = {}
Expand Down Expand Up @@ -389,6 +400,9 @@ async function evaluateWorkspace(options: CodyBenchOptions, recordingDirectory:
case BenchStrategy.Chat:
await evaluateChatStrategy(client, options)
break
case BenchStrategy.UnitTest:
await evaluateUnitTestStrategy(client, options)
break
default:
throw new Error(`unknown strategy ${options.fixture.strategy}`)
}
Expand All @@ -399,3 +413,30 @@ async function evaluateWorkspace(options: CodyBenchOptions, recordingDirectory:
await client.request('shutdown', null)
client.notify('exit', null)
}

function expandWorkspaces(
workspaces: CodyBenchOptions[] | undefined,
rootDir: string
): CodyBenchOptions[] {
if (!workspaces) {
return []
}
return workspaces.flatMap(workspace => {
workspace.absolutePath = path.normalize(path.join(rootDir, workspace.workspace))

if (!workspace.workspace.endsWith('/*')) {
return [workspace]
}
return glob
.sync(workspace.workspace, {
cwd: rootDir,
})
.flatMap(workspacePath => {
return {
...workspace,
workspace: workspacePath,
absolutePath: path.normalize(path.join(rootDir, workspacePath)),
}
})
})
}
Loading

0 comments on commit 6ebe1f1

Please sign in to comment.