Skip to content

Commit d2c9ff1

Browse files
leigaolleigaol
committed
add unit test
1 parent f75c58a commit d2c9ff1

File tree

2 files changed

+258
-55
lines changed

2 files changed

+258
-55
lines changed

‎packages/amazonq/test/unit/codewhisperer/tracker/userWrittenCodeTracker.test.ts‎

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import assert from 'assert'
7+
import * as sinon from 'sinon'
8+
import * as vscode from 'vscode'
9+
import { UserWrittenCodeTracker, TelemetryHelper, AuthUtil } from 'aws-core-vscode/codewhisperer'
10+
import { createMockDocument, resetCodeWhispererGlobalVariables } from 'aws-core-vscode/test'
11+
12+
describe('userWrittenCodeTracker', function () {
13+
describe('test isActive', function () {
14+
afterEach(async function () {
15+
await resetCodeWhispererGlobalVariables()
16+
UserWrittenCodeTracker.instance.reset()
17+
sinon.restore()
18+
})
19+
20+
it('inactive case: telemetryEnable = true, isConnected = false', function () {
21+
sinon.stub(TelemetryHelper.instance, 'isTelemetryEnabled').returns(true)
22+
sinon.stub(AuthUtil.instance, 'isConnected').returns(false)
23+
assert.strictEqual(UserWrittenCodeTracker.instance.isActive(), false)
24+
})
25+
26+
it('inactive case: telemetryEnabled = false, isConnected = false', function () {
27+
sinon.stub(TelemetryHelper.instance, 'isTelemetryEnabled').returns(false)
28+
sinon.stub(AuthUtil.instance, 'isConnected').returns(false)
29+
assert.strictEqual(UserWrittenCodeTracker.instance.isActive(), false)
30+
})
31+
32+
it('active case: telemetryEnabled = true, isConnected = true', function () {
33+
sinon.stub(TelemetryHelper.instance, 'isTelemetryEnabled').returns(true)
34+
sinon.stub(AuthUtil.instance, 'isConnected').returns(true)
35+
assert.strictEqual(UserWrittenCodeTracker.instance.isActive(), true)
36+
})
37+
})
38+
39+
describe('onDocumentChange', function () {
40+
let tracker: UserWrittenCodeTracker | undefined
41+
42+
beforeEach(async function () {
43+
await resetCodeWhispererGlobalVariables()
44+
tracker = UserWrittenCodeTracker.instance
45+
if (tracker) {
46+
sinon.stub(tracker, 'isActive').returns(true)
47+
}
48+
})
49+
50+
afterEach(function () {
51+
sinon.restore()
52+
UserWrittenCodeTracker.instance.reset()
53+
})
54+
55+
it('Should skip when content change size is more than 50', function () {
56+
if (!tracker) {
57+
assert.fail()
58+
}
59+
tracker.onQFeatureInvoked()
60+
tracker.onTextDocumentChange({
61+
reason: undefined,
62+
document: createMockDocument(),
63+
contentChanges: [
64+
{
65+
range: new vscode.Range(0, 0, 0, 600),
66+
rangeOffset: 0,
67+
rangeLength: 600,
68+
text: 'def twoSum(nums, target):\nfor '.repeat(20),
69+
},
70+
],
71+
})
72+
assert.strictEqual(tracker.getUserWrittenCharacters('python'), 0)
73+
assert.strictEqual(tracker.getUserWrittenLines('python'), 0)
74+
})
75+
76+
it('Should not skip when content change size is less than 50', function () {
77+
if (!tracker) {
78+
assert.fail()
79+
}
80+
tracker.onQFeatureInvoked()
81+
tracker.onTextDocumentChange({
82+
reason: undefined,
83+
document: createMockDocument(),
84+
contentChanges: [
85+
{
86+
range: new vscode.Range(0, 0, 0, 49),
87+
rangeOffset: 0,
88+
rangeLength: 49,
89+
text: 'a = 123'.repeat(7),
90+
},
91+
],
92+
})
93+
assert.strictEqual(tracker.getUserWrittenCharacters('python'), 0)
94+
assert.strictEqual(tracker.getUserWrittenLines('python'), 0)
95+
})
96+
97+
it('Should skip when CodeWhisperer is editing', function () {
98+
if (!tracker) {
99+
assert.fail()
100+
}
101+
tracker.onQFeatureInvoked()
102+
tracker.onQStartsMakingEdits()
103+
tracker.onTextDocumentChange({
104+
reason: undefined,
105+
document: createMockDocument(),
106+
contentChanges: [
107+
{
108+
range: new vscode.Range(0, 0, 0, 30),
109+
rangeOffset: 0,
110+
rangeLength: 30,
111+
text: 'def twoSum(nums, target):\nfor',
112+
},
113+
],
114+
})
115+
assert.strictEqual(tracker.getUserWrittenCharacters('python'), 0)
116+
assert.strictEqual(tracker.getUserWrittenLines('python'), 0)
117+
})
118+
119+
it('Should not reduce tokens when delete', function () {
120+
if (!tracker) {
121+
assert.fail()
122+
}
123+
const doc = createMockDocument('import math', 'test.py', 'python')
124+
125+
tracker.onQFeatureInvoked()
126+
tracker.onTextDocumentChange({
127+
reason: undefined,
128+
document: doc,
129+
contentChanges: [
130+
{
131+
range: new vscode.Range(0, 0, 0, 1),
132+
rangeOffset: 0,
133+
rangeLength: 0,
134+
text: 'a',
135+
},
136+
],
137+
})
138+
tracker.onTextDocumentChange({
139+
reason: undefined,
140+
document: doc,
141+
contentChanges: [
142+
{
143+
range: new vscode.Range(0, 0, 0, 1),
144+
rangeOffset: 0,
145+
rangeLength: 0,
146+
text: 'b',
147+
},
148+
],
149+
})
150+
assert.strictEqual(tracker.getUserWrittenCharacters('python'), 0)
151+
tracker.onTextDocumentChange({
152+
reason: undefined,
153+
document: doc,
154+
contentChanges: [
155+
{
156+
range: new vscode.Range(0, 0, 0, 1),
157+
rangeOffset: 1,
158+
rangeLength: 1,
159+
text: '',
160+
},
161+
],
162+
})
163+
assert.strictEqual(tracker.getUserWrittenCharacters('python'), 0)
164+
})
165+
})
166+
167+
describe('emitCodeWhispererCodeContribution', function () {
168+
let tracker: UserWrittenCodeTracker | undefined
169+
170+
beforeEach(async function () {
171+
await resetCodeWhispererGlobalVariables()
172+
tracker = UserWrittenCodeTracker.instance
173+
tracker.reset()
174+
if (tracker) {
175+
sinon.stub(tracker, 'isActive').returns(true)
176+
}
177+
})
178+
179+
afterEach(function () {
180+
sinon.restore()
181+
})
182+
183+
it('should emit correct code coverage telemetry in python file', async function () {})
184+
185+
it('Should not emit if user has not use any Q feature', async function () {})
186+
})
187+
})

‎packages/core/src/codewhisperer/tracker/userWrittenCodeTracker.ts‎

Lines changed: 71 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,16 @@ import { AuthUtil } from '../util/authUtil'
1111
import { getSelectedCustomization } from '../util/customizationUtil'
1212
import { codeWhispererClient as client } from '../client/codewhisperer'
1313
import { isAwsError } from '../../shared/errors'
14+
import { CodewhispererLanguage } from '../../shared'
1415

1516
/**
1617
* This singleton class is mainly used for calculating the user written code
1718
* for active Amazon Q users.
1819
* It reports the user written code per 5 minutes when the user is coding and using Amazon Q features
1920
*/
2021
export class UserWrittenCodeTracker {
21-
private _userWrittenNewCodeCharacterCount: number
22-
private _userWrittenNewCodeLineCount: number
22+
private _userWrittenNewCodeCharacterCount: Map<CodewhispererLanguage, number>
23+
private _userWrittenNewCodeLineCount: Map<CodewhispererLanguage, number>
2324
private _qIsMakingEdits: boolean
2425
private _timer?: NodeJS.Timer
2526
private _qUsageCount: number
@@ -30,22 +31,14 @@ export class UserWrittenCodeTracker {
3031
static resetQIsEditingTimeoutMs = 5 * 60 * 1000
3132
static defaultCheckPeriodMillis = 1000 * 60 * 5
3233
private constructor() {
33-
this._userWrittenNewCodeLineCount = 0
34-
this._userWrittenNewCodeCharacterCount = 0
34+
this._userWrittenNewCodeLineCount = new Map<CodewhispererLanguage, number>()
35+
this._userWrittenNewCodeCharacterCount = new Map<CodewhispererLanguage, number>()
3536
this._qUsageCount = 0
3637
this._qIsMakingEdits = false
3738
this._timer = undefined
3839
this._lastQInvocationTime = 0
3940
}
4041

41-
private resetTracker() {
42-
this._userWrittenNewCodeLineCount = 0
43-
this._userWrittenNewCodeCharacterCount = 0
44-
this._qUsageCount = 0
45-
this._qIsMakingEdits = false
46-
this._lastQInvocationTime = 0
47-
}
48-
4942
public static get instance() {
5043
return (this.#instance ??= new this())
5144
}
@@ -69,34 +62,61 @@ export class UserWrittenCodeTracker {
6962
this._qIsMakingEdits = false
7063
}
7164

72-
public emitCodeContribution() {
65+
public getUserWrittenCharacters(language: CodewhispererLanguage) {
66+
return this._userWrittenNewCodeCharacterCount.get(language) || 0
67+
}
68+
69+
public getUserWrittenLines(language: CodewhispererLanguage) {
70+
return this._userWrittenNewCodeLineCount.get(language) || 0
71+
}
72+
73+
public reset() {
74+
this._userWrittenNewCodeLineCount = new Map<CodewhispererLanguage, number>()
75+
this._userWrittenNewCodeCharacterCount = new Map<CodewhispererLanguage, number>()
76+
this._qUsageCount = 0
77+
this._qIsMakingEdits = false
78+
this._lastQInvocationTime = 0
79+
if (this._timer !== undefined) {
80+
clearTimeout(this._timer)
81+
this._timer = undefined
82+
}
83+
}
84+
85+
public emitCodeContributions() {
7386
const selectedCustomization = getSelectedCustomization()
74-
client
75-
.sendTelemetryEvent({
76-
telemetryEvent: {
77-
codeCoverageEvent: {
78-
customizationArn: selectedCustomization.arn === '' ? undefined : selectedCustomization.arn,
79-
programmingLanguage: {
80-
languageName: 'plaintext',
87+
88+
for (const [language, charCount] of this._userWrittenNewCodeCharacterCount) {
89+
const lineCount = this._userWrittenNewCodeLineCount.get(language) || 0
90+
if (charCount > 0) {
91+
client
92+
.sendTelemetryEvent({
93+
telemetryEvent: {
94+
codeCoverageEvent: {
95+
customizationArn:
96+
selectedCustomization.arn === '' ? undefined : selectedCustomization.arn,
97+
programmingLanguage: {
98+
languageName: language,
99+
},
100+
acceptedCharacterCount: 0,
101+
totalCharacterCount: 0,
102+
timestamp: new Date(Date.now()),
103+
userWrittenCodeCharacterCount: charCount,
104+
userWrittenCodeLineCount: lineCount,
105+
},
81106
},
82-
acceptedCharacterCount: 0,
83-
totalCharacterCount: 0,
84-
timestamp: new Date(Date.now()),
85-
userWrittenCodeCharacterCount: this._userWrittenNewCodeCharacterCount,
86-
userWrittenCodeLineCount: this._userWrittenNewCodeLineCount,
87-
},
88-
},
89-
})
90-
.then()
91-
.catch((error) => {
92-
let requestId: string | undefined
93-
if (isAwsError(error)) {
94-
requestId = error.requestId
95-
}
96-
getLogger().debug(
97-
`Failed to sendTelemetryEvent, requestId: ${requestId ?? ''}, message: ${error.message}`
98-
)
99-
})
107+
})
108+
.then()
109+
.catch((error) => {
110+
let requestId: string | undefined
111+
if (isAwsError(error)) {
112+
requestId = error.requestId
113+
}
114+
getLogger().debug(
115+
`Failed to sendTelemetryEvent, requestId: ${requestId ?? ''}, message: ${error.message}`
116+
)
117+
})
118+
}
119+
}
100120
}
101121

102122
private tryStartTimer() {
@@ -105,8 +125,7 @@ export class UserWrittenCodeTracker {
105125
}
106126
if (!this.isActive()) {
107127
getLogger().debug(`Skip emiting code contribution metric. Telemetry disabled or not logged in. `)
108-
this.resetTracker()
109-
this.closeTimer()
128+
this.reset()
110129
return
111130
}
112131
const startTime = performance.now()
@@ -120,28 +139,20 @@ export class UserWrittenCodeTracker {
120139
getLogger().debug(`Skip emiting code contribution metric. There is no active Amazon Q usage. `)
121140
return
122141
}
123-
if (this._userWrittenNewCodeCharacterCount === 0) {
142+
if (this._userWrittenNewCodeCharacterCount.size === 0) {
124143
getLogger().debug(`Skip emiting code contribution metric. There is no new code added. `)
125144
return
126145
}
127-
this.emitCodeContribution()
146+
this.emitCodeContributions()
128147
}
129148
} catch (e) {
130149
getLogger().verbose(`Exception Thrown from QCodeGenTracker: ${e}`)
131150
} finally {
132-
this.resetTracker()
133-
this.closeTimer()
151+
this.reset()
134152
}
135153
}, UserWrittenCodeTracker.defaultCheckPeriodMillis)
136154
}
137155

138-
private closeTimer() {
139-
if (this._timer !== undefined) {
140-
clearTimeout(this._timer)
141-
this._timer = undefined
142-
}
143-
}
144-
145156
private countNewLines(str: string) {
146157
return str.split('\n').length - 1
147158
}
@@ -170,9 +181,14 @@ export class UserWrittenCodeTracker {
170181
if (contentChange.text.length > UserWrittenCodeTracker.copySnippetThreshold) {
171182
return
172183
}
173-
this._userWrittenNewCodeCharacterCount += contentChange.text.length
174-
this._userWrittenNewCodeLineCount += this.countNewLines(contentChange.text)
175-
// start 5 min data reporting once valid user input is detected
176-
this.tryStartTimer()
184+
const language = runtimeLanguageContext.normalizeLanguage(e.document.languageId)
185+
if (language) {
186+
const charCount = this._userWrittenNewCodeCharacterCount.get(language) || 0
187+
this._userWrittenNewCodeCharacterCount.set(language, charCount + contentChange.text.length)
188+
const lineCount = this._userWrittenNewCodeLineCount.get(language) || 0
189+
this._userWrittenNewCodeLineCount.set(language, lineCount + this.countNewLines(contentChange.text))
190+
// start 5 min data reporting once valid user input is detected
191+
this.tryStartTimer()
192+
}
177193
}
178194
}

0 commit comments

Comments
 (0)