Skip to content

Commit 8cd9eda

Browse files
leigaolleigaol
committed
add unit test
1 parent 04a7dce commit 8cd9eda

File tree

3 files changed

+203
-5
lines changed

3 files changed

+203
-5
lines changed

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/UserWrittenCodeTracker.kt‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import com.intellij.psi.PsiDocumentManager
1313
import com.intellij.util.Alarm
1414
import com.intellij.util.AlarmFactory
1515
import com.intellij.util.messages.Topic
16+
import org.jetbrains.annotations.TestOnly
1617
import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
1718
import software.aws.toolkits.core.utils.debug
1819
import software.aws.toolkits.core.utils.getLogger
@@ -29,12 +30,12 @@ import java.util.concurrent.atomic.AtomicInteger
2930

3031
@Service(Service.Level.PROJECT)
3132
class UserWrittenCodeTracker(private val project: Project) : Disposable {
32-
private val userWrittenCodeLineCount = mutableMapOf<CodeWhispererProgrammingLanguage, Long>()
33-
private val userWrittenCodeCharacterCount = mutableMapOf<CodeWhispererProgrammingLanguage, Long>()
33+
val userWrittenCodeLineCount = mutableMapOf<CodeWhispererProgrammingLanguage, Long>()
34+
val userWrittenCodeCharacterCount = mutableMapOf<CodeWhispererProgrammingLanguage, Long>()
3435
private val alarm = AlarmFactory.getInstance().create(Alarm.ThreadToUse.POOLED_THREAD, this)
3536

3637
private val isShuttingDown = AtomicBoolean(false)
37-
private val qInvocationCount: AtomicInteger = AtomicInteger(0)
38+
val qInvocationCount: AtomicInteger = AtomicInteger(0)
3839
private val isQMakingEdits = AtomicBoolean(false)
3940
private val isActive: AtomicBoolean = AtomicBoolean(false)
4041

@@ -60,7 +61,7 @@ class UserWrittenCodeTracker(private val project: Project) : Disposable {
6061
scheduleTracker()
6162
}
6263

63-
private fun reset() {
64+
fun reset() {
6465
userWrittenCodeLineCount.clear()
6566
userWrittenCodeCharacterCount.clear()
6667
qInvocationCount.set(0)
@@ -171,6 +172,11 @@ class UserWrittenCodeTracker(private val project: Project) : Disposable {
171172
}
172173
flush()
173174
}
175+
176+
@TestOnly
177+
fun forceTrackerFlush() {
178+
alarm.drainRequestsInTest()
179+
}
174180
}
175181

176182
enum class QFeatureEvent {

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -391,7 +391,7 @@ class CodeWhispererClientAdaptorTest {
391391
@Test
392392
fun `sendTelemetryEvent for codePercentage respects telemetry optin status`() {
393393
sendTelemetryEventOptOutCheckHelper {
394-
sut.sendCodePercentageTelemetry(aProgrammingLanguage(), aString(), 0, 1, 0)
394+
sut.sendCodePercentageTelemetry(aProgrammingLanguage(), aString(), 0, 1, 0, 0, 0)
395395
}
396396
}
397397

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/UserWrittenCodeTrackerTest.kt‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.codewhisperer
5+
6+
import com.intellij.openapi.application.ApplicationManager
7+
import com.intellij.openapi.command.WriteCommandAction
8+
import com.intellij.openapi.editor.Editor
9+
import com.intellij.openapi.editor.event.DocumentEvent
10+
import com.intellij.openapi.project.Project
11+
import com.intellij.testFramework.DisposableRule
12+
import com.intellij.testFramework.fixtures.CodeInsightTestFixture
13+
import com.intellij.testFramework.replaceService
14+
import com.intellij.testFramework.runInEdtAndWait
15+
import org.assertj.core.api.Assertions.assertThat
16+
import org.junit.After
17+
import org.junit.Before
18+
import org.junit.Rule
19+
import org.junit.Test
20+
import org.mockito.internal.verification.Times
21+
import org.mockito.kotlin.any
22+
import org.mockito.kotlin.argumentCaptor
23+
import org.mockito.kotlin.doReturn
24+
import org.mockito.kotlin.mock
25+
import org.mockito.kotlin.spy
26+
import org.mockito.kotlin.verify
27+
import software.aws.toolkits.core.telemetry.TelemetryBatcher
28+
import software.aws.toolkits.core.telemetry.TelemetryPublisher
29+
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
30+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName
31+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
32+
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
33+
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
34+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
35+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython
36+
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.QFeatureEvent
37+
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.UserWrittenCodeTracker
38+
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.UserWrittenCodeTracker.Companion.Q_FEATURE_TOPIC
39+
import software.aws.toolkits.jetbrains.services.telemetry.NoOpPublisher
40+
import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
41+
import software.aws.toolkits.jetbrains.settings.AwsSettings
42+
import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
43+
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
44+
45+
internal class UserWrittenCodeTrackerTest {
46+
47+
protected class TestTelemetryService(
48+
publisher: TelemetryPublisher = NoOpPublisher(),
49+
batcher: TelemetryBatcher,
50+
) : TelemetryService(publisher, batcher)
51+
52+
@Rule
53+
@JvmField
54+
val projectRule: CodeInsightTestFixtureRule
55+
56+
@Rule
57+
@JvmField
58+
val disposableRule = DisposableRule()
59+
60+
@Rule
61+
@JvmField
62+
val mockClientManagerRule = MockClientManagerRule()
63+
64+
protected lateinit var project: Project
65+
protected lateinit var fixture: CodeInsightTestFixture
66+
protected lateinit var telemetryServiceSpy: TelemetryService
67+
protected lateinit var batcher: TelemetryBatcher
68+
protected lateinit var exploreActionManagerMock: CodeWhispererExplorerActionManager
69+
protected lateinit var sut: UserWrittenCodeTracker
70+
71+
init {
72+
this.projectRule = PythonCodeInsightTestFixtureRule()
73+
}
74+
75+
@Before
76+
open fun setup() {
77+
this.project = projectRule.project
78+
this.fixture = projectRule.fixture
79+
fixture.configureByText(pythonFileName, pythonTestLeftContext)
80+
AwsSettings.getInstance().isTelemetryEnabled = true
81+
batcher = mock()
82+
telemetryServiceSpy = spy(TestTelemetryService(batcher = batcher))
83+
84+
exploreActionManagerMock = mock {
85+
on { checkActiveCodeWhispererConnectionType(any()) } doReturn CodeWhispererLoginType.Sono
86+
}
87+
88+
ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, exploreActionManagerMock, disposableRule.disposable)
89+
ApplicationManager.getApplication().replaceService(TelemetryService::class.java, telemetryServiceSpy, disposableRule.disposable)
90+
91+
fixture.configureByText(pythonFileName, pythonTestLeftContext)
92+
runInEdtAndWait {
93+
projectRule.fixture.editor.caretModel.primaryCaret.moveToOffset(projectRule.fixture.editor.document.textLength)
94+
}
95+
}
96+
97+
@After
98+
fun tearDown() {
99+
if (::sut.isInitialized) {
100+
sut.forceTrackerFlush()
101+
sut.reset()
102+
}
103+
}
104+
105+
@Test
106+
fun `test tracker is listening to q service invocation`() {
107+
sut = UserWrittenCodeTracker.getInstance(project)
108+
sut.activateTrackerIfNotActive()
109+
assertThat(sut.qInvocationCount.get()).isEqualTo(0)
110+
ApplicationManager.getApplication().messageBus.syncPublisher(Q_FEATURE_TOPIC).onEvent(QFeatureEvent.INVOCATION)
111+
assertThat(sut.qInvocationCount.get()).isEqualTo(1)
112+
ApplicationManager.getApplication().messageBus.syncPublisher(Q_FEATURE_TOPIC).onEvent(QFeatureEvent.INVOCATION)
113+
assertThat(sut.qInvocationCount.get()).isEqualTo(2)
114+
}
115+
116+
117+
@Test
118+
fun `test tracker is not listening to multi char input more than 50, but works for less than 50, and will not increment totalTokens - add new code`() {
119+
sut = UserWrittenCodeTracker.getInstance(project)
120+
sut.activateTrackerIfNotActive()
121+
fixture.configureByText(pythonFileName, "")
122+
val newCode = "def addTwoNumbers\n return"
123+
runInEdtAndWait {
124+
WriteCommandAction.runWriteCommandAction(project) {
125+
fixture.editor.appendString(newCode)
126+
}
127+
}
128+
val language: CodeWhispererProgrammingLanguage = CodeWhispererPython.INSTANCE
129+
assertThat(sut.userWrittenCodeCharacterCount[language]).isEqualTo(newCode.length.toLong())
130+
assertThat(sut.userWrittenCodeLineCount[language]).isEqualTo(1)
131+
132+
val anotherCode = "(x, y):\n".repeat(8)
133+
runInEdtAndWait {
134+
WriteCommandAction.runWriteCommandAction(project) {
135+
fixture.editor.appendString(anotherCode)
136+
}
137+
}
138+
assertThat(sut.userWrittenCodeCharacterCount[language]).isEqualTo(newCode.length.toLong())
139+
assertThat(sut.userWrittenCodeLineCount[language]).isEqualTo(1)
140+
}
141+
142+
@Test
143+
fun `test tracker is listening to document changes and increment totalTokens - delete code should not affect`() {
144+
sut = UserWrittenCodeTracker.getInstance(project)
145+
sut.activateTrackerIfNotActive()
146+
assertThat(sut.userWrittenCodeCharacterCount.getOrDefault(CodeWhispererPython.INSTANCE, 0)).isEqualTo(0)
147+
runInEdtAndWait {
148+
fixture.editor.caretModel.primaryCaret.moveToOffset(fixture.editor.document.textLength)
149+
WriteCommandAction.runWriteCommandAction(project) {
150+
fixture.editor.document.deleteString(fixture.editor.caretModel.offset - 3, fixture.editor.caretModel.offset)
151+
}
152+
}
153+
assertThat(sut.userWrittenCodeCharacterCount.getOrDefault(CodeWhispererPython.INSTANCE, 0)).isEqualTo(0)
154+
}
155+
156+
157+
@Test
158+
fun `test tracker is listening to document changes only when Q is not editing`() {
159+
sut = UserWrittenCodeTracker.getInstance(project)
160+
sut.activateTrackerIfNotActive()
161+
fixture.configureByText(pythonFileName, "")
162+
val newCode = "def addTwoNumbers\n return"
163+
164+
ApplicationManager.getApplication().messageBus.syncPublisher(Q_FEATURE_TOPIC).onEvent(QFeatureEvent.STARTS_EDITING)
165+
runInEdtAndWait {
166+
WriteCommandAction.runWriteCommandAction(project) {
167+
fixture.editor.appendString(newCode)
168+
}
169+
}
170+
171+
ApplicationManager.getApplication().messageBus.syncPublisher(Q_FEATURE_TOPIC).onEvent(QFeatureEvent.FINISHES_EDITING)
172+
val language: CodeWhispererProgrammingLanguage = CodeWhispererPython.INSTANCE
173+
assertThat(sut.userWrittenCodeCharacterCount.getOrDefault(language, 0)).isEqualTo(0)
174+
assertThat(sut.userWrittenCodeLineCount.getOrDefault(language, 0)).isEqualTo(0)
175+
176+
runInEdtAndWait {
177+
WriteCommandAction.runWriteCommandAction(project) {
178+
fixture.editor.appendString(newCode)
179+
}
180+
}
181+
assertThat(sut.userWrittenCodeCharacterCount[CodeWhispererPython.INSTANCE]).isEqualTo(newCode.length.toLong())
182+
assertThat(sut.userWrittenCodeLineCount[CodeWhispererPython.INSTANCE]).isEqualTo(1)
183+
184+
}
185+
186+
private fun Editor.appendString(string: String) {
187+
val currentOffset = caretModel.primaryCaret.offset
188+
document.insertString(currentOffset, string)
189+
caretModel.moveToOffset(currentOffset + string.length)
190+
}
191+
192+
}

0 commit comments

Comments
 (0)