modified InjectHeader to also modify currentRequest (#245)

Co-authored-by: Fernando Draghi <fdraghi@FDraghi-Embrace-MBP.local>
embrace-io · May 14, 2024 · b8ef540 · b8ef540
1 parent 646b536
commit b8ef540
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/Sources/EmbraceObjCUtils/source/NSURLSessionTask+Embrace.m b/Sources/EmbraceObjCUtils/source/NSURLSessionTask+Embrace.m
@@ -12,13 +12,17 @@ - (BOOL)injectHeaderWithKey:(NSString *)key value:(NSString *)value {
         return NO;
     }
 
-    if (![self.originalRequest isKindOfClass:[NSMutableURLRequest class]]) {
+    if (![self.originalRequest isKindOfClass:[NSMutableURLRequest class]] ||
+        ![self.currentRequest isKindOfClass:[NSMutableURLRequest class]]) {
         return NO;
     }
 
     NSMutableURLRequest *request = (NSMutableURLRequest *)self.originalRequest;
     [request setValue:value forHTTPHeaderField:key];
 
+    NSMutableURLRequest *currentRequest = (NSMutableURLRequest *)self.currentRequest;
+    [currentRequest setValue:value forHTTPHeaderField:key];
+
     return YES;
 }
 

diff --git a/Tests/EmbraceCoreTests/Capture/Network/DefaultURLSessionTaskHandlerTests.swift b/Tests/EmbraceCoreTests/Capture/Network/DefaultURLSessionTaskHandlerTests.swift
@@ -392,6 +392,12 @@ private extension DefaultURLSessionTaskHandlerTests {
             let tracingHeader = headers!["traceparent"]
             XCTAssertNotNil(tracingHeader)
 
+            let currentHeaders = task.currentRequest?.allHTTPHeaderFields
+            XCTAssertNotNil(currentHeaders)
+
+            let currentTracingHeader = currentHeaders!["traceparent"]
+            XCTAssertNotNil(currentTracingHeader)
+
             let span = try XCTUnwrap(otel.spanProcessor.startedSpans.first)
             validateTracingHeaderForSpan(tracingHeader: tracingHeader!, span: span)
         } catch let exception {