Skip to content

Commit 46c48ef

Browse files
Will-ShaoHuaWill-ShaoHua
authored
[CodeWhisperer] %code tracker enhancement (#3293)
* cwspr %code tracker enhancement: will exclude documentChanged triggered by code reformat * restructure codewhispererCoverageTrackerTest
1 parent 7a3f724 commit 46c48ef

File tree

3 files changed

+137
-36
lines changed

3 files changed

+137
-36
lines changed

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorListener.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33

44
package software.aws.toolkits.jetbrains.services.codewhisperer.editor
55

6+
import com.intellij.openapi.editor.event.BulkAwareDocumentListener
67
import com.intellij.openapi.editor.event.DocumentEvent
7-
import com.intellij.openapi.editor.event.DocumentListener
88
import com.intellij.openapi.editor.event.EditorFactoryEvent
99
import com.intellij.openapi.editor.event.EditorFactoryListener
1010
import com.intellij.openapi.editor.impl.EditorImpl
@@ -25,7 +25,7 @@ class CodeWhispererEditorListener : EditorFactoryListener {
2525
if (!CodeWhispererLanguageManager.getInstance().isLanguageSupported(language)) return
2626
// If language is supported, install document listener for CodeWhisperer service
2727
editor.document.addDocumentListener(
28-
object : DocumentListener {
28+
object : BulkAwareDocumentListener {
2929
override fun documentChanged(event: DocumentEvent) {
3030
if (!CodeWhispererExplorerActionManager.getInstance().hasAcceptedTermsOfService()) return
3131
CodeWhispererInvocationStatus.getInstance().documentChanged()

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererCodeCoverageTracker.kt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,8 @@ abstract class CodeWhispererCodeCoverageTracker(
9191
LOG.debug { "event with isWholeTextReplaced flag: $event" }
9292
if (event.oldTimeStamp == 0L) return
9393
}
94+
// This case capture IDE reformatting the document, which will be blank string
95+
if (isDocumentEventFromReformatting(event)) return
9496
incrementTotalTokens(event.document, event.newLength - event.oldLength)
9597
}
9698

@@ -147,6 +149,9 @@ abstract class CodeWhispererCodeCoverageTracker(
147149
}
148150
}
149151

152+
private fun isDocumentEventFromReformatting(event: DocumentEvent): Boolean =
153+
(event.newFragment.toString().isBlank() && event.oldFragment.toString().isBlank()) && (event.oldLength == 0 || event.newLength == 0)
154+
150155
private fun reset() {
151156
startTime = Instant.now()
152157
rangeMarkers.clear()

‎jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt‎

Lines changed: 130 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import com.intellij.openapi.editor.RangeMarker
1111
import com.intellij.openapi.editor.event.DocumentEvent
1212
import com.intellij.openapi.project.Project
1313
import com.intellij.openapi.util.Key
14+
import com.intellij.psi.codeStyle.CodeStyleManager
1415
import com.intellij.testFramework.DisposableRule
1516
import com.intellij.testFramework.fixtures.CodeInsightTestFixture
1617
import com.intellij.testFramework.replaceService
@@ -56,54 +57,50 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhisperer
5657
import software.aws.toolkits.jetbrains.services.telemetry.NoOpPublisher
5758
import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
5859
import software.aws.toolkits.jetbrains.settings.AwsSettings
60+
import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
61+
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
5962
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
6063
import software.aws.toolkits.telemetry.CodewhispererCompletionType
6164
import software.aws.toolkits.telemetry.CodewhispererLanguage
6265

63-
class CodeWhispererCodeCoverageTrackerTest {
64-
private class TestCodePercentageTracker(
66+
internal abstract class CodeWhispererCodeCoverageTrackerTestBase(myProjectRule: CodeInsightTestFixtureRule) {
67+
protected class TestCodePercentageTracker(
6568
timeWindowInSec: Long,
6669
language: CodewhispererLanguage,
6770
rangeMarkers: MutableList<RangeMarker> = mutableListOf(),
6871
codeCoverageTokens: MutableMap<Document, CodeCoverageTokens> = mutableMapOf()
6972
) : CodeWhispererCodeCoverageTracker(timeWindowInSec, language, rangeMarkers, codeCoverageTokens)
7073

71-
private class TestTelemetryService(
74+
protected class TestTelemetryService(
7275
publisher: TelemetryPublisher = NoOpPublisher(),
7376
batcher: TelemetryBatcher
7477
) : TelemetryService(publisher, batcher)
75-
7678
@Rule
7779
@JvmField
78-
var projectRule = PythonCodeInsightTestFixtureRule()
80+
val projectRule: CodeInsightTestFixtureRule
7981

8082
@Rule
8183
@JvmField
82-
val mockClientManagerRule = MockClientManagerRule()
84+
val disposableRule = DisposableRule()
8385

8486
@Rule
8587
@JvmField
86-
val disposableRule = DisposableRule()
88+
val mockClientManagerRule = MockClientManagerRule()
8789

88-
private lateinit var project: Project
89-
private lateinit var fixture: CodeInsightTestFixture
90-
private lateinit var telemetryServiceSpy: TelemetryService
91-
private lateinit var batcher: TelemetryBatcher
92-
private lateinit var exploreActionManagerMock: CodeWhispererExplorerActionManager
90+
protected lateinit var project: Project
91+
protected lateinit var fixture: CodeInsightTestFixture
92+
protected lateinit var telemetryServiceSpy: TelemetryService
93+
protected lateinit var batcher: TelemetryBatcher
94+
protected lateinit var exploreActionManagerMock: CodeWhispererExplorerActionManager
9395

94-
private lateinit var invocationContext: InvocationContext
95-
private lateinit var sessionContext: SessionContext
96+
init {
97+
this.projectRule = myProjectRule
98+
}
9699

97-
@Before
98-
fun setup() {
100+
open fun setup() {
101+
this.project = projectRule.project
102+
this.fixture = projectRule.fixture
99103
AwsSettings.getInstance().isTelemetryEnabled = true
100-
project = projectRule.project
101-
fixture = projectRule.fixture
102-
fixture.configureByText(pythonFileName, pythonTestLeftContext)
103-
runInEdtAndWait {
104-
projectRule.fixture.editor.caretModel.primaryCaret.moveToOffset(projectRule.fixture.editor.document.textLength)
105-
}
106-
107104
batcher = mock()
108105
telemetryServiceSpy = spy(TestTelemetryService(batcher = batcher))
109106
exploreActionManagerMock = mock {
@@ -112,7 +109,32 @@ class CodeWhispererCodeCoverageTrackerTest {
112109

113110
ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, exploreActionManagerMock, disposableRule.disposable)
114111
ApplicationManager.getApplication().replaceService(TelemetryService::class.java, telemetryServiceSpy, disposableRule.disposable)
115-
project.replaceService(CodeWhispererCodeReferenceManager::class.java, mock(), disposableRule.disposable)
112+
}
113+
114+
@After
115+
fun tearDown() {
116+
CodeWhispererCodeCoverageTracker.getInstancesMap().clear()
117+
}
118+
119+
protected companion object {
120+
const val CODE_PERCENTAGE = "codewhisperer_codePercentage"
121+
const val CWSPR_PERCENTAGE = "codewhispererPercentage"
122+
const val CWSPR_Language = "codewhispererLanguage"
123+
const val CWSPR_ACCEPTED_TOKENS = "codewhispererAcceptedTokens"
124+
const val CWSPR_TOTAL_TOKENS = "codewhispererTotalTokens"
125+
}
126+
}
127+
128+
internal class CodeWhispererCodeCoverageTrackerTestPython : CodeWhispererCodeCoverageTrackerTestBase(PythonCodeInsightTestFixtureRule()) {
129+
private lateinit var invocationContext: InvocationContext
130+
private lateinit var sessionContext: SessionContext
131+
@Before
132+
override fun setup() {
133+
super.setup()
134+
fixture.configureByText(pythonFileName, pythonTestLeftContext)
135+
runInEdtAndWait {
136+
projectRule.fixture.editor.caretModel.primaryCaret.moveToOffset(projectRule.fixture.editor.document.textLength)
137+
}
116138

117139
val requestContext = RequestContext(
118140
project,
@@ -131,11 +153,9 @@ class CodeWhispererCodeCoverageTrackerTest {
131153
)
132154
invocationContext = InvocationContext(requestContext, responseContext, recommendationContext, mock())
133155
sessionContext = SessionContext()
134-
}
135156

136-
@After
137-
fun tearDown() {
138-
CodeWhispererCodeCoverageTracker.getInstancesMap().clear()
157+
// it is needed because referenceManager is listening to CODEWHISPERER_USER_ACTION_PERFORMED topic
158+
project.replaceService(CodeWhispererCodeReferenceManager::class.java, mock(), disposableRule.disposable)
139159
}
140160

141161
@Test
@@ -207,6 +227,26 @@ class CodeWhispererCodeCoverageTrackerTest {
207227
assertThat(pythonTracker.totalTokensSize).isEqualTo(pythonTestLeftContext.length - 3)
208228
}
209229

230+
@Test
231+
fun `test tracker documentChanged - will not increment tokens on blank string of length greater than 1`() {
232+
val pythonTracker = TestCodePercentageTracker(
233+
TOTAL_SECONDS_IN_MINUTE,
234+
CodewhispererLanguage.Python,
235+
codeCoverageTokens = mutableMapOf(fixture.editor.document to CodeCoverageTokens(pythonTestLeftContext.length, 0))
236+
)
237+
CodeWhispererCodeCoverageTracker.getInstancesMap()[CodewhispererLanguage.Python] = pythonTracker
238+
pythonTracker.activateTrackerIfNotActive()
239+
assertThat(pythonTracker.totalTokensSize).isEqualTo(pythonTestLeftContext.length)
240+
241+
runInEdtAndWait {
242+
WriteCommandAction.runWriteCommandAction(project) {
243+
fixture.editor.document.insertString(fixture.editor.caretModel.offset, "\t")
244+
}
245+
}
246+
247+
assertThat(pythonTracker.totalTokensSize).isEqualTo(pythonTestLeftContext.length)
248+
}
249+
210250
@Test
211251
fun `test msg CODEWHISPERER_USER_ACTION_PERFORMED will add rangeMarker in the list`() {
212252
val pythonTracker = spy(TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, language = CodewhispererLanguage.Python))
@@ -400,12 +440,68 @@ class CodeWhispererCodeCoverageTrackerTest {
400440
document.insertString(currentOffset, string)
401441
caretModel.moveToOffset(currentOffset + string.length)
402442
}
443+
}
403444

404-
private companion object {
405-
const val CODE_PERCENTAGE = "codewhisperer_codePercentage"
406-
const val CWSPR_PERCENTAGE = "codewhispererPercentage"
407-
const val CWSPR_Language = "codewhispererLanguage"
408-
const val CWSPR_ACCEPTED_TOKENS = "codewhispererAcceptedTokens"
409-
const val CWSPR_TOTAL_TOKENS = "codewhispererTotalTokens"
445+
internal class CodeWhispererCodeCoverageTrackerTestJava : CodeWhispererCodeCoverageTrackerTestBase(JavaCodeInsightTestFixtureRule()) {
446+
@Before
447+
override fun setup() {
448+
super.setup()
449+
}
450+
451+
@Test
452+
fun `tracker should not update totalTokens if documentChanged events are fired by code reformatting`() {
453+
val codeNeedToBeReformatted = """
454+
class Answer {
455+
private int knapsack(int[] w, int[] v, int c) {
456+
int[][] dp = new int[w.length + 1][c + 1];
457+
for (int i = 0; i < w.length; i++) {for (int j = 0; j <= c; j++) {
458+
if (j < w[i]) {
459+
dp[i + 1][j] = dp[i][j];
460+
}
461+
else {
462+
dp[i + 1][j] = Math.max(dp[i][j], dp[i][j - w[i]] + v[i]);
463+
}
464+
}
465+
}
466+
return dp[w.length][c];
467+
}
468+
}
469+
""".trimIndent()
470+
val file = fixture.configureByText("test.java", codeNeedToBeReformatted)
471+
val tracker = spy(
472+
TestCodePercentageTracker(
473+
TOTAL_SECONDS_IN_MINUTE,
474+
language = CodewhispererLanguage.Java,
475+
codeCoverageTokens = mutableMapOf(fixture.editor.document to CodeCoverageTokens(totalTokens = codeNeedToBeReformatted.length))
476+
)
477+
)
478+
CodeWhispererCodeCoverageTracker.getInstancesMap()[CodewhispererLanguage.Java] = tracker
479+
runInEdtAndWait {
480+
WriteCommandAction.runWriteCommandAction(project) {
481+
CodeStyleManager.getInstance(project).reformatText(file, 0, fixture.editor.document.textLength)
482+
}
483+
}
484+
// reformat should fire documentChanged events, but tracker should not update token from these events
485+
verify(tracker, atLeastOnce()).documentChanged(any())
486+
assertThat(tracker.totalTokensSize).isEqualTo(codeNeedToBeReformatted.length)
487+
488+
val formatted = """
489+
class Answer {
490+
private int knapsack(int[] w, int[] v, int c) {
491+
int[][] dp = new int[w.length + 1][c + 1];
492+
for (int i = 0; i < w.length; i++) {
493+
for (int j = 0; j <= c; j++) {
494+
if (j < w[i]) {
495+
dp[i + 1][j] = dp[i][j];
496+
} else {
497+
dp[i + 1][j] = Math.max(dp[i][j], dp[i][j - w[i]] + v[i]);
498+
}
499+
}
500+
}
501+
return dp[w.length][c];
502+
}
503+
}
504+
""".trimIndent()
505+
assertThat(fixture.editor.document.text.trimEnd()).isEqualTo(formatted)
410506
}
411507
}

0 commit comments

Comments
 (0)