@@ -2,8 +2,10 @@ package com.datadog.appsec.api.security
2
2
3
3
import com.datadog.appsec.event.EventProducerService
4
4
import com.datadog.appsec.event.ExpiredSubscriberInfoException
5
+ import com.datadog.appsec.event.data.KnownAddresses
5
6
import com.datadog.appsec.gateway.AppSecRequestContext
6
7
import datadog.trace.api.gateway.RequestContext
8
+ import datadog.trace.api.internal.TraceSegment
7
9
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
8
10
import datadog.trace.test.util.DDSpecification
9
11
@@ -16,6 +18,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
16
18
def subInfo = Mock (EventProducerService.DataSubscriberInfo )
17
19
def span = Mock (AgentSpan )
18
20
def reqCtx = Mock (RequestContext )
21
+ def traceSegment = Mock (TraceSegment )
19
22
def ctx = Mock (AppSecRequestContext )
20
23
def processor = new AppSecSpanPostProcessor (sampler, producer)
21
24
@@ -28,9 +31,11 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
28
31
1 * reqCtx. getData(_) >> ctx
29
32
1 * ctx. isKeepOpenForApiSecurityPostProcessing() >> true
30
33
1 * sampler. sampleRequest(_) >> true
31
- 1 * producer. getDataSubscribers(_) >> subInfo
34
+ 1 * reqCtx. getTraceSegment() >> traceSegment
35
+ 1 * producer. getDataSubscribers(KnownAddresses . WAF_CONTEXT_PROCESSOR ) >> subInfo
32
36
1 * subInfo. isEmpty() >> false
33
37
1 * producer. publishDataEvent(_, ctx, _, _)
38
+ 1 * ctx. commitDerivatives(traceSegment)
34
39
1 * ctx. setKeepOpenForApiSecurityPostProcessing(false )
35
40
1 * ctx. close()
36
41
1 * sampler. releaseOne()
@@ -67,6 +72,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
67
72
def producer = Mock (EventProducerService )
68
73
def span = Mock (AgentSpan )
69
74
def reqCtx = Mock (RequestContext )
75
+ def traceSegment = Mock (TraceSegment )
70
76
def ctx = Mock (AppSecRequestContext )
71
77
def processor = new AppSecSpanPostProcessor (sampler, producer)
72
78
@@ -79,6 +85,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
79
85
1 * reqCtx. getData(_) >> ctx
80
86
1 * ctx. isKeepOpenForApiSecurityPostProcessing() >> true
81
87
1 * sampler. sampleRequest(_) >> true
88
+ 1 * reqCtx. getTraceSegment() >> traceSegment
82
89
1 * producer. getDataSubscribers(_) >> null
83
90
1 * ctx. setKeepOpenForApiSecurityPostProcessing(false )
84
91
1 * ctx. close() >> { throw new RuntimeException () }
@@ -184,6 +191,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
184
191
def subInfo = Mock (EventProducerService.DataSubscriberInfo )
185
192
def span = Mock (AgentSpan )
186
193
def reqCtx = Mock (RequestContext )
194
+ def traceSegment = Mock (TraceSegment )
187
195
def ctx = Mock (AppSecRequestContext )
188
196
def processor = new AppSecSpanPostProcessor (sampler, producer)
189
197
@@ -196,6 +204,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
196
204
1 * reqCtx. getData(_) >> ctx
197
205
1 * ctx. isKeepOpenForApiSecurityPostProcessing() >> true
198
206
1 * sampler. sampleRequest(_) >> true
207
+ 1 * reqCtx. getTraceSegment() >> traceSegment
199
208
1 * producer. getDataSubscribers(_) >> subInfo
200
209
1 * subInfo. isEmpty() >> true
201
210
1 * ctx. setKeepOpenForApiSecurityPostProcessing(false )
@@ -211,6 +220,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
211
220
def subInfo = Mock (EventProducerService.DataSubscriberInfo )
212
221
def span = Mock (AgentSpan )
213
222
def reqCtx = Mock (RequestContext )
223
+ def traceSegment = Mock (TraceSegment )
214
224
def ctx = Mock (AppSecRequestContext )
215
225
def processor = new AppSecSpanPostProcessor (sampler, producer)
216
226
@@ -223,6 +233,7 @@ class AppSecSpanPostProcessorTest extends DDSpecification {
223
233
1 * reqCtx. getData(_) >> ctx
224
234
1 * ctx. isKeepOpenForApiSecurityPostProcessing() >> true
225
235
1 * sampler. sampleRequest(_) >> true
236
+ 1 * reqCtx. getTraceSegment() >> traceSegment
226
237
1 * producer. getDataSubscribers(_) >> subInfo
227
238
1 * subInfo. isEmpty() >> false
228
239
1 * producer. publishDataEvent(_, ctx, _, _) >> { throw new ExpiredSubscriberInfoException () }
0 commit comments