@@ -2344,4 +2344,94 @@ mod tests {
23442344 "no contexts should be ready to send yet"
23452345 ) ;
23462346 }
2347+
2348+ fn setup_appsec ( ) -> Processor {
2349+ let aws_config = Arc :: new ( AwsConfig {
2350+ region : "us-east-1" . into ( ) ,
2351+ aws_lwa_proxy_lambda_runtime_api : Some ( "***" . into ( ) ) ,
2352+ function_name : "test-function" . into ( ) ,
2353+ sandbox_init_time : Instant :: now ( ) ,
2354+ runtime_api : "***" . into ( ) ,
2355+ exec_wrapper : None ,
2356+ initialization_type : "on-demand" . into ( ) ,
2357+ } ) ;
2358+ let config = Arc :: new ( config:: Config {
2359+ service : Some ( "test-service" . to_string ( ) ) ,
2360+ serverless_appsec_enabled : true ,
2361+ ..config:: Config :: default ( )
2362+ } ) ;
2363+ let tags_provider = Arc :: new ( provider:: Provider :: new (
2364+ Arc :: clone ( & config) ,
2365+ LAMBDA_RUNTIME_SLUG . to_string ( ) ,
2366+ & HashMap :: from ( [ ( "function_arn" . to_string ( ) , "test-arn" . to_string ( ) ) ] ) ,
2367+ ) ) ;
2368+ let ( service, handle) =
2369+ dogstatsd:: aggregator:: AggregatorService :: new ( dogstatsd:: metric:: EMPTY_TAGS , 1024 )
2370+ . expect ( "failed to create aggregator service" ) ;
2371+ tokio:: spawn ( service. run ( ) ) ;
2372+ let propagator = Arc :: new ( DatadogCompositePropagator :: new ( Arc :: clone ( & config) ) ) ;
2373+ let ( durable_context_tx, _) = tokio:: sync:: mpsc:: channel ( 1 ) ;
2374+ Processor :: new ( tags_provider, config, aws_config, handle, propagator, durable_context_tx)
2375+ }
2376+
2377+ #[ tokio:: test]
2378+ async fn enrich_ctx_sets_appsec_enabled_when_aap_enabled ( ) {
2379+ let mut p = setup_appsec ( ) ;
2380+ let request_id = String :: from ( "req-appsec" ) ;
2381+ p. on_invoke_event ( request_id. clone ( ) ) ;
2382+ p. on_platform_start ( request_id. clone ( ) , chrono:: Utc :: now ( ) ) ;
2383+
2384+ let ctx = p
2385+ . enrich_ctx_at_platform_done ( & request_id, Status :: Success )
2386+ . expect ( "context must be present" ) ;
2387+
2388+ assert_eq ! (
2389+ ctx. invocation_span. metrics. get( "_dd.appsec.enabled" ) ,
2390+ Some ( & 1.0 ) ,
2391+ "_dd.appsec.enabled must be 1.0 when AAP is enabled"
2392+ ) ;
2393+ }
2394+
2395+ #[ tokio:: test]
2396+ async fn enrich_ctx_does_not_set_appsec_enabled_when_aap_disabled ( ) {
2397+ let mut p = setup ( ) ;
2398+ let request_id = String :: from ( "req-no-appsec" ) ;
2399+ p. on_invoke_event ( request_id. clone ( ) ) ;
2400+ p. on_platform_start ( request_id. clone ( ) , chrono:: Utc :: now ( ) ) ;
2401+
2402+ let ctx = p
2403+ . enrich_ctx_at_platform_done ( & request_id, Status :: Success )
2404+ . expect ( "context must be present" ) ;
2405+
2406+ assert ! (
2407+ ctx. invocation_span. metrics. get( "_dd.appsec.enabled" ) . is_none( ) ,
2408+ "_dd.appsec.enabled must not be set when AAP is disabled"
2409+ ) ;
2410+ }
2411+
2412+ #[ tokio:: test]
2413+ async fn enrich_ctx_does_not_override_existing_appsec_enabled ( ) {
2414+ let mut p = setup_appsec ( ) ;
2415+ let request_id = String :: from ( "req-appsec-preset" ) ;
2416+ p. on_invoke_event ( request_id. clone ( ) ) ;
2417+ p. on_platform_start ( request_id. clone ( ) , chrono:: Utc :: now ( ) ) ;
2418+
2419+ // Pre-set a different value to verify or_insert does not overwrite it.
2420+ p. context_buffer
2421+ . get_mut ( & request_id)
2422+ . expect ( "context must exist" )
2423+ . invocation_span
2424+ . metrics
2425+ . insert ( "_dd.appsec.enabled" . to_string ( ) , 0.0 ) ;
2426+
2427+ let ctx = p
2428+ . enrich_ctx_at_platform_done ( & request_id, Status :: Success )
2429+ . expect ( "context must be present" ) ;
2430+
2431+ assert_eq ! (
2432+ ctx. invocation_span. metrics. get( "_dd.appsec.enabled" ) ,
2433+ Some ( & 0.0 ) ,
2434+ "pre-existing _dd.appsec.enabled value must not be overwritten"
2435+ ) ;
2436+ }
23472437}
0 commit comments