Skip to content

Commit 841947a

Browse files
samgst-amazonsamgst-amazon
committed
handle telemetry/event messages
1 parent f24037f commit 841947a

File tree

2 files changed

+265
-1
lines changed

2 files changed

+265
-1
lines changed

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,20 +11,55 @@ import org.eclipse.lsp4j.MessageParams
1111
import org.eclipse.lsp4j.MessageType
1212
import org.eclipse.lsp4j.PublishDiagnosticsParams
1313
import org.eclipse.lsp4j.ShowMessageRequestParams
14+
import software.amazon.awssdk.services.toolkittelemetry.model.MetricUnit
15+
import software.aws.toolkits.core.utils.getLogger
16+
import software.aws.toolkits.core.utils.warn
1417
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
1518
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
1619
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
1720
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.ConnectionMetadata
1821
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.SsoProfileData
22+
import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
1923
import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
24+
import java.time.Instant
2025
import java.util.concurrent.CompletableFuture
2126

2227
/**
2328
* Concrete implementation of [AmazonQLanguageClient] to handle messages sent from server
2429
*/
2530
class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageClient {
2631
override fun telemetryEvent(`object`: Any) {
27-
println(`object`)
32+
when (`object`) {
33+
is MutableMap<*, *> -> handleTelemetryMap(`object`)
34+
}
35+
}
36+
37+
private fun handleTelemetryMap(telemetryMap: MutableMap<*, *>) {
38+
try {
39+
val name = telemetryMap["name"] as? String ?: return
40+
41+
@Suppress("UNCHECKED_CAST")
42+
val data = telemetryMap["data"] as? MutableMap<String, Any> ?: return
43+
44+
TelemetryService.getInstance().record(project) {
45+
datum(name) {
46+
createTime(Instant.now())
47+
unit(telemetryMap["unit"] as? MetricUnit ?: MetricUnit.NONE)
48+
value(telemetryMap["value"] as? Double ?: 1.0)
49+
passive(telemetryMap["passive"] as? Boolean ?: false)
50+
51+
telemetryMap["result"]?.let { result ->
52+
metadata("result", result.toString())
53+
}
54+
55+
data.forEach { (key, value) ->
56+
metadata(key, value.toString())
57+
}
58+
}
59+
}
60+
} catch (e: Exception) {
61+
LOG.warn(e) { "Failed to process telemetry event: $telemetryMap" }
62+
}
2863
}
2964

3065
override fun publishDiagnostics(diagnostics: PublishDiagnosticsParams) {
@@ -92,4 +127,8 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
92127
}
93128
)
94129
}
130+
131+
companion object {
132+
private val LOG = getLogger<AmazonQLanguageClientImpl>()
133+
}
95134
}

‎plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImplTest.kt‎

Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,23 +10,248 @@ import com.intellij.openapi.project.Project
1010
import com.intellij.testFramework.ApplicationExtension
1111
import io.mockk.every
1212
import io.mockk.mockk
13+
import io.mockk.mockkObject
14+
import io.mockk.slot
15+
import io.mockk.verify
1316
import org.assertj.core.api.Assertions.assertThat
17+
import org.assertj.core.api.Assertions.entry
1418
import org.eclipse.lsp4j.ConfigurationItem
1519
import org.eclipse.lsp4j.ConfigurationParams
1620
import org.junit.jupiter.api.Test
1721
import org.junit.jupiter.api.extension.ExtendWith
22+
import software.amazon.awssdk.services.toolkittelemetry.model.MetricUnit
23+
import software.aws.toolkits.core.telemetry.DefaultMetricEvent
24+
import software.aws.toolkits.core.telemetry.MetricEvent
1825
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
1926
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
2027
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
2128
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.ConnectionMetadata
2229
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.SsoProfileData
30+
import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
2331
import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
2432

2533
@ExtendWith(ApplicationExtension::class)
2634
class AmazonQLanguageClientImplTest {
2735
private val project: Project = mockk(relaxed = true)
2836
private val sut = AmazonQLanguageClientImpl(project)
2937

38+
@Test
39+
fun `telemetryEvent handles basic event with name and data`() {
40+
val telemetryService = mockk<TelemetryService>(relaxed = true)
41+
mockkObject(TelemetryService)
42+
every { TelemetryService.getInstance() } returns telemetryService
43+
44+
val builderCaptor = slot<MetricEvent.Builder.() -> Unit>()
45+
every { telemetryService.record(project, capture(builderCaptor)) } returns Unit
46+
47+
val event = mapOf(
48+
"name" to "test_event",
49+
"data" to mapOf(
50+
"key1" to "value1",
51+
"key2" to 42
52+
)
53+
)
54+
55+
sut.telemetryEvent(event)
56+
57+
val builder = DefaultMetricEvent.builder()
58+
builderCaptor.captured.invoke(builder)
59+
60+
val metricEvent = builder.build()
61+
val datum = metricEvent.data.first()
62+
63+
assertThat(datum.name).isEqualTo("test_event")
64+
assertThat(datum.metadata).contains(
65+
entry("key1", "value1"),
66+
entry("key2", "42")
67+
)
68+
}
69+
70+
@Test
71+
fun `telemetryEvent handles event with result field`() {
72+
val telemetryService = mockk<TelemetryService>(relaxed = true)
73+
mockkObject(TelemetryService)
74+
every { TelemetryService.getInstance() } returns telemetryService
75+
76+
val builderCaptor = slot<MetricEvent.Builder.() -> Unit>()
77+
every { telemetryService.record(project, capture(builderCaptor)) } returns Unit
78+
79+
val event = mapOf(
80+
"name" to "test_event",
81+
"result" to "success",
82+
"data" to mapOf(
83+
"key1" to "value1"
84+
)
85+
)
86+
87+
sut.telemetryEvent(event)
88+
89+
val builder = DefaultMetricEvent.builder()
90+
builderCaptor.captured.invoke(builder)
91+
92+
val metricEvent = builder.build()
93+
val datum = metricEvent.data.first()
94+
95+
assertThat(datum.name).isEqualTo("test_event")
96+
assertThat(datum.metadata).contains(
97+
entry("key1", "value1"),
98+
entry("result", "success")
99+
)
100+
}
101+
102+
@Test
103+
fun `telemetryEvent uses custom unit value when provided`() {
104+
val telemetryService = mockk<TelemetryService>(relaxed = true)
105+
mockkObject(TelemetryService)
106+
every { TelemetryService.getInstance() } returns telemetryService
107+
108+
val builderCaptor = slot<MetricEvent.Builder.() -> Unit>()
109+
every { telemetryService.record(project, capture(builderCaptor)) } returns Unit
110+
111+
val event = mapOf(
112+
"name" to "test_event",
113+
"unit" to MetricUnit.MILLISECONDS,
114+
"data" to mapOf(
115+
"key1" to "value1"
116+
)
117+
)
118+
119+
sut.telemetryEvent(event)
120+
121+
val builder = DefaultMetricEvent.builder()
122+
builderCaptor.captured.invoke(builder)
123+
124+
val metricEvent = builder.build()
125+
val datum = metricEvent.data.first()
126+
127+
assertThat(datum.unit).isEqualTo(MetricUnit.MILLISECONDS)
128+
assertThat(datum.metadata).contains(entry("key1", "value1"))
129+
}
130+
131+
@Test
132+
fun `telemetryEvent uses custom value when provided`() {
133+
val telemetryService = mockk<TelemetryService>(relaxed = true)
134+
mockkObject(TelemetryService)
135+
every { TelemetryService.getInstance() } returns telemetryService
136+
137+
val builderCaptor = slot<MetricEvent.Builder.() -> Unit>()
138+
every { telemetryService.record(project, capture(builderCaptor)) } returns Unit
139+
140+
val event = mapOf(
141+
"name" to "test_event",
142+
"value" to 2.5,
143+
"data" to mapOf(
144+
"key1" to "value1"
145+
)
146+
)
147+
148+
sut.telemetryEvent(event)
149+
150+
val builder = DefaultMetricEvent.builder()
151+
builderCaptor.captured.invoke(builder)
152+
153+
val metricEvent = builder.build()
154+
val datum = metricEvent.data.first()
155+
156+
assertThat(datum.value).isEqualTo(2.5)
157+
assertThat(datum.metadata).contains(entry("key1", "value1"))
158+
}
159+
160+
@Test
161+
fun `telemetryEvent uses custom passive value when provided`() {
162+
val telemetryService = mockk<TelemetryService>(relaxed = true)
163+
mockkObject(TelemetryService)
164+
every { TelemetryService.getInstance() } returns telemetryService
165+
166+
val builderCaptor = slot<MetricEvent.Builder.() -> Unit>()
167+
every { telemetryService.record(project, capture(builderCaptor)) } returns Unit
168+
169+
val event = mapOf(
170+
"name" to "test_event",
171+
"passive" to true,
172+
"data" to mapOf(
173+
"key1" to "value1"
174+
)
175+
)
176+
177+
sut.telemetryEvent(event)
178+
179+
val builder = DefaultMetricEvent.builder()
180+
builderCaptor.captured.invoke(builder)
181+
182+
val metricEvent = builder.build()
183+
val datum = metricEvent.data.first()
184+
185+
assertThat(datum.passive).isTrue()
186+
assertThat(datum.metadata).contains(entry("key1", "value1"))
187+
}
188+
189+
@Test
190+
fun `telemetryEvent ignores event without name`() {
191+
val telemetryService = mockk<TelemetryService>(relaxed = true)
192+
mockkObject(TelemetryService)
193+
every { TelemetryService.getInstance() } returns telemetryService
194+
195+
val event = mapOf(
196+
"data" to mapOf(
197+
"key1" to "value1"
198+
)
199+
)
200+
201+
sut.telemetryEvent(event)
202+
203+
verify(exactly = 0) {
204+
telemetryService.record(project, any())
205+
}
206+
}
207+
208+
@Test
209+
fun `telemetryEvent ignores event without data`() {
210+
val telemetryService = mockk<TelemetryService>(relaxed = true)
211+
mockkObject(TelemetryService)
212+
every { TelemetryService.getInstance() } returns telemetryService
213+
214+
val event = mapOf(
215+
"name" to "test_event"
216+
)
217+
218+
sut.telemetryEvent(event)
219+
220+
verify(exactly = 0) {
221+
telemetryService.record(project, any())
222+
}
223+
}
224+
225+
@Test
226+
fun `telemetryEvent uses default values when not provided`() {
227+
val telemetryService = mockk<TelemetryService>(relaxed = true)
228+
mockkObject(TelemetryService)
229+
every { TelemetryService.getInstance() } returns telemetryService
230+
231+
val builderCaptor = slot<MetricEvent.Builder.() -> Unit>()
232+
every { telemetryService.record(project, capture(builderCaptor)) } returns Unit
233+
234+
val event = mapOf(
235+
"name" to "test_event",
236+
"data" to mapOf(
237+
"key1" to "value1"
238+
)
239+
)
240+
241+
sut.telemetryEvent(event)
242+
243+
val builder = DefaultMetricEvent.builder()
244+
builderCaptor.captured.invoke(builder)
245+
246+
val metricEvent = builder.build()
247+
val datum = metricEvent.data.first()
248+
249+
assertThat(datum.unit).isEqualTo(MetricUnit.NONE)
250+
assertThat(datum.value).isEqualTo(1.0)
251+
assertThat(datum.passive).isFalse()
252+
assertThat(datum.metadata).contains(entry("key1", "value1"))
253+
}
254+
30255
@Test
31256
fun `getConnectionMetadata returns connection metadata with start URL for bearer token connection`() {
32257
val mockConnectionManager = mockk<ToolkitConnectionManager>()

0 commit comments

Comments
 (0)