Forward all write events in injecting NIO handler

Author: slashmo

Sources/NIOInstrumentation/HTTPHeadersInjectingHandler.swift renamed to Sources/NIOInstrumentation/HeaderInjectingHTTPClientHandler.swift

@@ -16,15 +16,18 @@ import Instrumentation
 import NIO
 import NIOHTTP1
 
-public final class HTTPHeadersInjectingHandler: ChannelOutboundHandler {
+public final class HeaderInjectingHTTPClientHandler: ChannelOutboundHandler {
     public typealias OutboundIn = HTTPClientRequestPart
     public typealias OutboundOut = HTTPClientRequestPart
 
     public init() {}
 
     public func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
-        let requestPart = unwrapOutboundIn(data)
-        guard case .head(var head) = requestPart else { return }
+        guard case .head(var head) = self.unwrapOutboundIn(data) else {
+            context.write(data, promise: promise)
+            return
+        }
+
         InstrumentationSystem.instrument.inject(context.baggage, into: &head.headers, using: HTTPHeadersInjector())
         context.write(self.wrapOutboundOut(.head(head)), promise: promise)
     }

Tests/LinuxMain.swift

@@ -37,8 +37,8 @@ class LinuxMainRunnerImpl: LinuxMainRunner {
        XCTMain([
              testCase(HTTPHeadersCarrierTests.allTests),
              testCase(HTTPHeadersExtractInjectTests.allTests),
-             testCase(HTTPHeadersInjectingHandlerTests.allTests),
              testCase(HeaderExtractingHTTPServerHandlerTests.allTests),
+             testCase(HeaderInjectingHTTPClientHandlerTests.allTests),
              testCase(InstrumentTests.allTests),
              testCase(InstrumentationSystemTests.allTests),
              testCase(SpanAttributeSemanticsTests.allTests),

Tests/NIOInstrumentationTests/HTTPHeadersExtractInjectTests.swift

@@ -26,7 +26,7 @@ final class HTTPHeadersExtractInjectTests: XCTestCase {
     func test_extracted_baggage_into_subsequent_request_headers() throws {
         InstrumentationSystem.bootstrapInternal(FakeTracer())
 
-        let outboundHandler = HTTPHeadersInjectingHandler()
+        let outboundHandler = HeaderInjectingHTTPClientHandler()
         let requestHandler = MockRequestHandler()
         let inboundHandler = HeaderExtractingHTTPServerHandler()
 

Tests/NIOInstrumentationTests/HeaderExtractingHTTPServerHandlerTests.swift

@@ -61,20 +61,23 @@ final class HeaderExtractingHTTPServerHandlerTests: XCTestCase {
     }
 
     func test_forwards_all_read_events() throws {
-        InstrumentationSystem.bootstrapInternal(FakeTracer())
-
-        let channel = EmbeddedChannel(loop: EmbeddedEventLoop())
-        try channel.pipeline.addHandlers(HeaderExtractingHTTPServerHandler()).wait()
+        let channel = EmbeddedChannel(
+            handler: HeaderExtractingHTTPServerHandler(),
+            loop: EmbeddedEventLoop()
+        )
 
         let requestHead = HTTPRequestHead(version: .init(major: 1, minor: 1), method: .GET, uri: "/")
-        try channel.writeInbound(HTTPServerRequestPart.head(requestHead))
-        XCTAssertNotNil(try channel.readInbound(as: HTTPServerRequestPart.self))
+        let head = HTTPServerRequestPart.head(requestHead)
+        try channel.writeInbound(head)
+        XCTAssertEqual(try channel.readInbound(), head)
 
-        try channel.writeInbound(HTTPServerRequestPart.body(channel.allocator.buffer(string: "Test")))
-        XCTAssertNotNil(try channel.readInbound(as: HTTPServerRequestPart.self))
+        let body = HTTPServerRequestPart.body(channel.allocator.buffer(string: "Test"))
+        try channel.writeInbound(body)
+        XCTAssertEqual(try channel.readInbound(), body)
 
-        try channel.writeInbound(HTTPServerRequestPart.end(nil))
-        XCTAssertNotNil(try channel.readInbound(as: HTTPServerRequestPart.self))
+        let end = HTTPServerRequestPart.end(nil)
+        try channel.writeInbound(end)
+        XCTAssertEqual(try channel.readInbound(as: HTTPServerRequestPart.self), end)
     }
 }
 

Tests/NIOInstrumentationTests/HTTPHeadersInjectingHandlerTests+XCTest.swift renamed to Tests/NIOInstrumentationTests/HeaderInjectingHTTPClientHandlerTests+XCTest.swift

@@ -11,7 +11,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// HTTPHeadersInjectingHandlerTests+XCTest.swift
+// HeaderInjectingHTTPClientHandlerTests+XCTest.swift
 //
 import XCTest
 ///
@@ -20,12 +20,13 @@ import XCTest
 /// Do NOT edit this file directly as it will be regenerated automatically when needed.
 ///
 
-extension HTTPHeadersInjectingHandlerTests {
+extension HeaderInjectingHTTPClientHandlerTests {
 
    @available(*, deprecated, message: "not actually deprecated. Just deprecated to allow deprecated tests (which test deprecated functionality) without warnings")
-   static var allTests : [(String, (HTTPHeadersInjectingHandlerTests) -> () throws -> Void)] {
+   static var allTests : [(String, (HeaderInjectingHTTPClientHandlerTests) -> () throws -> Void)] {
       return [
                 ("test_injects_baggage_into_http_request_headers", test_injects_baggage_into_http_request_headers),
+                ("test_forwards_all_write_events", test_forwards_all_write_events),
            ]
    }
 }

Tests/NIOInstrumentationTests/HTTPHeadersInjectingHandlerTests.swift renamed to Tests/NIOInstrumentationTests/HeaderInjectingHTTPClientHandlerTests.swift

@@ -18,7 +18,9 @@ import NIOHTTP1
 import NIOInstrumentation
 import XCTest
 
-final class HTTPHeadersInjectingHandlerTests: XCTestCase {
+final class HeaderInjectingHTTPClientHandlerTests: XCTestCase {
+    private let httpVersion = HTTPVersion(major: 1, minor: 1)
+
     override class func tearDown() {
         super.tearDown()
         InstrumentationSystem.bootstrapInternal(nil)
@@ -32,19 +34,37 @@ final class HTTPHeadersInjectingHandlerTests: XCTestCase {
         var baggage = Baggage.topLevel
         baggage[FakeTracer.TraceIDKey.self] = traceID
 
-        let httpVersion = HTTPVersion(major: 1, minor: 1)
-        let handler = HTTPHeadersInjectingHandler()
+        let handler = HeaderInjectingHTTPClientHandler()
         let loop = EmbeddedEventLoop()
         let channel = EmbeddedChannel(handler: handler, loop: loop)
         channel._channelCore.baggage = baggage
         let requestHead = HTTPRequestHead(version: httpVersion, method: .GET, uri: "/")
 
         try channel.writeOutbound(HTTPClientRequestPart.head(requestHead))
-        let modifiedRequestPart = try channel.readOutbound(as: HTTPClientRequestPart.self)
 
         XCTAssertEqual(
-            modifiedRequestPart,
-            .head(.init(version: httpVersion, method: .GET, uri: "/", headers: [FakeTracer.headerName: traceID]))
+            try channel.readOutbound(as: HTTPClientRequestPart.self),
+            .head(.init(version: self.httpVersion, method: .GET, uri: "/", headers: [FakeTracer.headerName: traceID]))
+        )
+    }
+
+    func test_forwards_all_write_events() throws {
+        let channel = EmbeddedChannel(
+            handler: HeaderInjectingHTTPClientHandler(),
+            loop: EmbeddedEventLoop()
         )
+
+        let requestHead = HTTPRequestHead(version: httpVersion, method: .GET, uri: "/")
+        let head = HTTPClientRequestPart.head(requestHead)
+        try channel.writeOutbound(head)
+        XCTAssertEqual(try channel.readOutbound(), head)
+
+        let body = HTTPClientRequestPart.body(.byteBuffer(channel.allocator.buffer(string: "test")))
+        try channel.writeOutbound(body)
+        XCTAssertEqual(try channel.readOutbound(), body)
+
+        let end = HTTPClientRequestPart.end(nil)
+        try channel.writeOutbound(end)
+        XCTAssertEqual(try channel.readOutbound(), end)
     }
 }