RUMM-1579 Add assertions for x-datadog-sampling-priority and `x-dat…

…adog-sampled` headers in integration tests
DataDog · Aug 27, 2021 · 1546f37 · 1546f37
1 parent 2e76012
commit 1546f37
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 23 deletions.
diff --git a/Tests/DatadogIntegrationTests/Scenarios/RUM/RUMResourcesScenarioTests.swift b/Tests/DatadogIntegrationTests/Scenarios/RUM/RUMResourcesScenarioTests.swift
@@ -98,10 +98,17 @@ class RUMResourcesScenarioTests: IntegrationTests, RUMCommonAsserts {
             getSpanID(from: firstPartyPOSTRequest),
             "Tracing information should be propagated to `firstPartyPOSTResourceURL`."
         )
-        XCTAssertEqual(
-            getDatadogOrigin(from: firstPartyPOSTRequest),
-            "rum",
-            "Additional `x-datadog-origin: rum` header should be propagated to `firstPartyPOSTResourceURL`"
+        XCTAssertTrue(
+            firstPartyPOSTRequest.httpHeaders.contains("x-datadog-sampling-priority: 1"),
+            "`x-datadog-sampling-priority: 1` header must be set for `firstPartyPOSTResourceURL`"
+        )
+        XCTAssertTrue(
+            firstPartyPOSTRequest.httpHeaders.contains("x-datadog-sampled: 1"),
+            "`x-datadog-sampled: 1` header must be set for `firstPartyPOSTResourceURL`"
+        )
+        XCTAssertTrue(
+            firstPartyPOSTRequest.httpHeaders.contains("x-datadog-origin: rum"),
+            "`x-datadog-origin: rum` header must be set for `firstPartyPOSTResourceURL`"
         )
 
         // Get RUM Sessions with expected number of View visits and Resources
@@ -232,11 +239,4 @@ class RUMResourcesScenarioTests: IntegrationTests, RUMCommonAsserts {
         header?.removeFirst(prefix.count)
         return header
     }
-
-    private func getDatadogOrigin(from request: Request) -> String? {
-        let prefix = "x-datadog-origin: "
-        var header = request.httpHeaders.first { $0.hasPrefix(prefix) }
-        header?.removeFirst(prefix.count)
-        return header
-    }
 }
diff --git a/Tests/DatadogIntegrationTests/Scenarios/Tracing/TracingURLSessionScenarioTests.swift b/Tests/DatadogIntegrationTests/Scenarios/Tracing/TracingURLSessionScenarioTests.swift
@@ -124,19 +124,27 @@ class TracingURLSessionScenarioTests: IntegrationTests, TracingCommonAsserts {
         XCTAssertEqual(firstPartyRequests.count, 1)
 
         let firstPartyRequest = firstPartyRequests[0]
-        let expectedTraceIDHeader = "x-datadog-trace-id: \(try taskWithRequest.traceID().hexadecimalNumberToDecimal)"
-        let expectedSpanIDHeader = "x-datadog-parent-id: \(try taskWithRequest.spanID().hexadecimalNumberToDecimal)"
-        XCTAssertTrue(
-            firstPartyRequest.httpHeaders.contains(expectedTraceIDHeader),
-            """
-            Request `\(firstPartyRequest.path)` does not contain `\(expectedTraceIDHeader)` header.
-            - request.headers: \(firstPartyRequest.httpHeaders)
-            """
-        )
-        XCTAssertTrue(
-            firstPartyRequest.httpHeaders.contains(expectedSpanIDHeader),
+        let expectedHeaders = [
+            "x-datadog-trace-id: \(try taskWithRequest.traceID().hexadecimalNumberToDecimal)",
+            "x-datadog-parent-id: \(try taskWithRequest.spanID().hexadecimalNumberToDecimal)",
+            "x-datadog-sampling-priority: 1",
+            "x-datadog-sampled: 1"
+        ]
+
+        expectedHeaders.forEach { expectedHeader in
+            XCTAssertTrue(
+                firstPartyRequest.httpHeaders.contains(expectedHeader),
+                """
+                Request `\(firstPartyRequest.path)` must contain `\(expectedHeader)` header.
+                - request.headers: \(firstPartyRequest.httpHeaders)
+                """
+            )
+        }
+
+        XCTAssertFalse(
+            firstPartyRequest.httpHeaders.contains("x-datadog-origin: rum"),
             """
-            Request `\(firstPartyRequest.path)` does not contain `\(expectedSpanIDHeader)` header.
+            Request `\(firstPartyRequest.path)` must not contain `x-datadog-origin: rum` header.
             - request.headers: \(firstPartyRequest.httpHeaders)
             """
         )