Skip to content

Commit a6e4f7d

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Fix coverage and telemetry tests
1 parent 0a9bdb9 commit a6e4f7d

File tree

3 files changed

+99
-16
lines changed

3 files changed

+99
-16
lines changed

‎dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -204,7 +204,7 @@ class AIGuardInternalTests extends DDSpecification {
204204
eval.action == suite.action
205205
eval.reason == suite.reason
206206
}
207-
assertTelemetry("ai_guard.requests", "action:$suite.action", "block:$throwAbortError")
207+
assertTelemetry('ai_guard.requests', "action:$suite.action", "block:$throwAbortError", 'error:false')
208208

209209
where:
210210
suite << TestSuite.build()
@@ -234,7 +234,7 @@ class AIGuardInternalTests extends DDSpecification {
234234
final exception = thrown(AIGuard.AIGuardClientError)
235235
exception.errors == errors
236236
1 * span.addThrowable(_ as AIGuard.AIGuardClientError)
237-
assertTelemetry("ai_guard.requests", "error:true")
237+
assertTelemetry('ai_guard.requests', 'error:true')
238238
}
239239

240240
void 'test evaluate with invalid JSON'() {
@@ -259,7 +259,7 @@ class AIGuardInternalTests extends DDSpecification {
259259
then:
260260
thrown(AIGuard.AIGuardClientError)
261261
1 * span.addThrowable(_ as AIGuard.AIGuardClientError)
262-
assertTelemetry("ai_guard.requests", "error:true")
262+
assertTelemetry('ai_guard.requests', 'error:true')
263263
}
264264

265265
void 'test evaluate with missing action'() {
@@ -284,7 +284,7 @@ class AIGuardInternalTests extends DDSpecification {
284284
then:
285285
thrown(AIGuard.AIGuardClientError)
286286
1 * span.addThrowable(_ as AIGuard.AIGuardClientError)
287-
assertTelemetry("ai_guard.requests", "error:true")
287+
assertTelemetry('ai_guard.requests', 'error:true')
288288
}
289289

290290
void 'test evaluate with non JSON response'() {
@@ -309,7 +309,7 @@ class AIGuardInternalTests extends DDSpecification {
309309
then:
310310
thrown(AIGuard.AIGuardClientError)
311311
1 * span.addThrowable(_ as AIGuard.AIGuardClientError)
312-
assertTelemetry("ai_guard.requests", "error:true")
312+
assertTelemetry('ai_guard.requests', 'error:true')
313313
}
314314

315315
void 'test evaluate with empty response'() {
@@ -334,7 +334,7 @@ class AIGuardInternalTests extends DDSpecification {
334334
then:
335335
thrown(AIGuard.AIGuardClientError)
336336
1 * span.addThrowable(_ as AIGuard.AIGuardClientError)
337-
assertTelemetry("ai_guard.requests", "error:true")
337+
assertTelemetry('ai_guard.requests', 'error:true')
338338
}
339339

340340
void 'test message length truncation'() {
@@ -366,7 +366,7 @@ class AIGuardInternalTests extends DDSpecification {
366366
assert received.size() == maxMessages
367367
assert received.size() < messages.size()
368368
}
369-
assertTelemetry("ai_guard.truncated", "type:messages")
369+
assertTelemetry('ai_guard.truncated', 'type:messages')
370370
}
371371

372372
void 'test message content truncation'() {
@@ -398,7 +398,7 @@ class AIGuardInternalTests extends DDSpecification {
398398
assert it.content.length() < message.content.length()
399399
}
400400
}
401-
assertTelemetry("ai_guard.truncated", "type:content")
401+
assertTelemetry('ai_guard.truncated', 'type:content')
402402
}
403403

404404
void 'test no messages'() {

‎internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -387,13 +387,13 @@ public void prepareMetrics() {
387387
for (final AIGuard.Action action : AIGuard.Action.values()) {
388388
final long blocked = aiGuardRequests.getAndSet(action.ordinal() * 2 + 1, 0);
389389
if (blocked > 0) {
390-
if (!rawMetricsQueue.offer(new AIGuardRequests(blocked, action, true))) {
390+
if (!rawMetricsQueue.offer(AIGuardRequests.success(blocked, action, true))) {
391391
break;
392392
}
393393
}
394394
final long nonBlocked = aiGuardRequests.getAndSet(action.ordinal() * 2, 0);
395395
if (nonBlocked > 0) {
396-
if (!rawMetricsQueue.offer(new AIGuardRequests(nonBlocked, action, false))) {
396+
if (!rawMetricsQueue.offer(AIGuardRequests.success(nonBlocked, action, false))) {
397397
break;
398398
}
399399
}
@@ -402,7 +402,7 @@ public void prepareMetrics() {
402402
// AI Guard failed requests
403403
final int aiGuardErrorRequests = aiGuardErrors.getAndSet(0);
404404
if (aiGuardErrorRequests > 0) {
405-
if (!rawMetricsQueue.offer(new AIGuardRequests(aiGuardErrorRequests, true))) {
405+
if (!rawMetricsQueue.offer(AIGuardRequests.error(aiGuardErrorRequests))) {
406406
return;
407407
}
408408
}
@@ -632,12 +632,17 @@ public WafInputTruncated(final long counter, final int bitfield) {
632632
}
633633

634634
public static class AIGuardRequests extends WafMetric {
635-
public AIGuardRequests(final long count, final AIGuard.Action action, final boolean block) {
636-
super("ai_guard.requests", count, "action:" + action, "block:" + block);
635+
private AIGuardRequests(final long count, final String... tags) {
636+
super("ai_guard.requests", count, tags);
637637
}
638638

639-
public AIGuardRequests(final long count, final boolean error) {
640-
super("ai_guard.requests", count, "error:" + error);
639+
public static AIGuardRequests success(
640+
final long count, final AIGuard.Action action, final boolean block) {
641+
return new AIGuardRequests(count, "action:" + action, "block:" + block, "error:false");
642+
}
643+
644+
public static AIGuardRequests error(final long count) {
645+
return new AIGuardRequests(count, "error:true");
641646
}
642647
}
643648

@@ -652,7 +657,7 @@ public enum AIGuardTruncationType {
652657
CONTENT("content");
653658
public final String tagValue;
654659

655-
AIGuardTruncationType(String tagValue) {
660+
AIGuardTruncationType(final String tagValue) {
656661
this.tagValue = tagValue;
657662
}
658663
}

‎internal-api/src/test/groovy/datadog/trace/api/telemetry/WafMetricCollectorTest.groovy‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,11 @@
11
package datadog.trace.api.telemetry
22

3+
import static datadog.trace.api.aiguard.AIGuard.Action.ABORT
4+
import static datadog.trace.api.aiguard.AIGuard.Action.ALLOW
5+
import static datadog.trace.api.aiguard.AIGuard.Action.DENY
6+
import static datadog.trace.api.telemetry.WafMetricCollector.AIGuardTruncationType.CONTENT
7+
import static datadog.trace.api.telemetry.WafMetricCollector.AIGuardTruncationType.MESSAGES
8+
39
import datadog.trace.test.util.DDSpecification
410

511
import java.util.concurrent.CountDownLatch
@@ -507,6 +513,78 @@ class WafMetricCollectorTest extends DDSpecification {
507513
metric.tags.toSet() == ['waf_version:waf_ver1', 'event_rules_version:rules.1'].toSet()
508514
}
509515

516+
void 'test ai guard request'() {
517+
given:
518+
final collector = WafMetricCollector.get()
519+
520+
when:
521+
collector.aiGuardRequest(action, block)
522+
523+
then:
524+
collector.prepareMetrics()
525+
final metrics = collector.drain()
526+
final configErrorMetrics = metrics.findAll { it.metricName == 'ai_guard.requests' }
527+
528+
final metric = configErrorMetrics[0]
529+
metric.type == 'count'
530+
metric.metricName == 'ai_guard.requests'
531+
metric.namespace == 'appsec'
532+
metric.value == 1
533+
metric.tags.toSet() == ["action:${action.name()}", "block:${block}", 'error:false'].toSet()
534+
535+
where:
536+
action | block
537+
ALLOW | true
538+
ALLOW | false
539+
DENY | true
540+
DENY | false
541+
ABORT | true
542+
ABORT | false
543+
}
544+
545+
void 'test ai guard error'() {
546+
given:
547+
final collector = WafMetricCollector.get()
548+
549+
when:
550+
collector.aiGuardError()
551+
552+
then:
553+
collector.prepareMetrics()
554+
final metrics = collector.drain()
555+
final configErrorMetrics = metrics.findAll { it.metricName == 'ai_guard.requests' }
556+
557+
final metric = configErrorMetrics[0]
558+
metric.type == 'count'
559+
metric.metricName == 'ai_guard.requests'
560+
metric.namespace == 'appsec'
561+
metric.value == 1
562+
metric.tags.toSet() == ['error:true'].toSet()
563+
}
564+
565+
void 'test ai guard truncated'() {
566+
given:
567+
final collector = WafMetricCollector.get()
568+
569+
when:
570+
collector.aiGuardTruncated(type)
571+
572+
then:
573+
collector.prepareMetrics()
574+
final metrics = collector.drain()
575+
final configErrorMetrics = metrics.findAll { it.metricName == 'ai_guard.truncated' }
576+
577+
final metric = configErrorMetrics[0]
578+
metric.type == 'count'
579+
metric.metricName == 'ai_guard.truncated'
580+
metric.namespace == 'appsec'
581+
metric.value == 1
582+
metric.tags.toSet() == ["type:${type.tagValue}"].toSet()
583+
584+
where:
585+
type << [MESSAGES, CONTENT]
586+
}
587+
510588
/**
511589
* Helper method to generate all combinations of n boolean values.
512590
*/

0 commit comments

Comments
 (0)