Fix flaky test

grpc · Feb 3, 2025 · 41b6e20 · 41b6e20
1 parent b549df5
commit 41b6e20
Showing 1 changed file with 40 additions and 84 deletions.
diff --git a/Tests/GRPCNIOTransportHTTP2Tests/HTTP2TransportTLSEnabledTests.swift b/Tests/GRPCNIOTransportHTTP2Tests/HTTP2TransportTLSEnabledTests.swift
@@ -154,8 +154,7 @@ struct HTTP2TransportTLSEnabledTests {
 
   @Test(
     "Error is surfaced when server fails client verification",
-    arguments: TransportKind.supported,
-    TransportKind.supported
+    arguments: [TransportKind.posix], [TransportKind.posix]
   )
   // Verification should fail because the client does not offer a cert that
   // the server can use for mutual verification.
@@ -229,8 +228,6 @@ struct HTTP2TransportTLSEnabledTests {
   enum TLSEnabledTestsError: Error {
     case failedToImportPKCS12
     case unexpectedListeningAddress
-    case serverError(cause: any Error)
-    case clientError(cause: any Error)
   }
 
   struct Config<Transport, Security> {
@@ -477,98 +474,57 @@ struct HTTP2TransportTLSEnabledTests {
     serverConfig: ServerConfig,
     _ test: (ControlClient<NIOClientTransport>) async throws -> Void
   ) async throws {
-    try await withThrowingDiscardingTaskGroup { group in
-      let server = self.makeServer(config: serverConfig)
-
-      group.addTask {
-        do {
-          try await server.serve()
-        } catch {
-          throw TLSEnabledTestsError.serverError(cause: error)
-        }
-      }
+    let serverTransport: NIOServerTransport
+    switch serverConfig {
+    case .posix(let posix):
+      serverTransport = NIOServerTransport(
+        .http2NIOPosix(
+          address: .ipv4(host: "127.0.0.1", port: 0),
+          transportSecurity: posix.security,
+          config: posix.transport
+        )
+      )
+    #if canImport(Network)
+    case .transportServices(let config):
+      serverTransport = NIOServerTransport(
+        .http2NIOTS(
+          address: .ipv4(host: "127.0.0.1", port: 0),
+          transportSecurity: config.security,
+          config: config.transport
+        )
+      )
+    #endif
+    }
 
+    try await withGRPCServer(transport: serverTransport, services: [ControlService()]) { server in
       guard let address = try await server.listeningAddress?.ipv4 else {
         throw TLSEnabledTestsError.unexpectedListeningAddress
       }
 
       let target: any ResolvableTarget = .ipv4(host: address.host, port: address.port)
-      let client = try self.makeClient(config: clientConfig, target: target)
-
-      group.addTask {
-        do {
-          try await client.runConnections()
-        } catch {
-          throw TLSEnabledTestsError.clientError(cause: error)
-        }
-      }
-
-      let control = ControlClient(wrapping: client)
-      try await test(control)
-
-      client.beginGracefulShutdown()
-      server.beginGracefulShutdown()
-    }
-  }
-
-  private func makeServer(config: ServerConfig) -> GRPCServer<NIOServerTransport> {
-    let services = [ControlService()]
-
-    switch config {
-    case .posix(let config):
-      return GRPCServer(
-        transport: NIOServerTransport(
+      let clientTransport: NIOClientTransport
+      switch clientConfig {
+      case .posix(let config):
+        clientTransport = try NIOClientTransport(
           .http2NIOPosix(
-            address: .ipv4(host: "127.0.0.1", port: 0),
+            target: target,
             transportSecurity: config.security,
             config: config.transport
           )
-        ),
-        services: services
-      )
-
-    #if canImport(Network)
-    case .transportServices(let config):
-      return GRPCServer(
-        transport: NIOServerTransport(
-          .http2NIOTS(
-            address: .ipv4(host: "127.0.0.1", port: 0),
-            transportSecurity: config.security,
-            config: config.transport
-          )
-        ),
-        services: services
-      )
-    #endif
-    }
-  }
-
-  private func makeClient(
-    config: ClientConfig,
-    target: any ResolvableTarget
-  ) throws -> GRPCClient<NIOClientTransport> {
-    switch config {
-    case .posix(let config):
-      let transport = try HTTP2ClientTransport.Posix(
-        target: target,
-        transportSecurity: config.security,
-        config: config.transport,
-        serviceConfig: ServiceConfig()
-      )
-      return GRPCClient(transport: NIOClientTransport(transport))
+        )
+      #if canImport(Network)
+      case .transportServices(let config):
+        clientTransport = try NIOClientTransport(
+          .http2NIOTS(target: target, transportSecurity: config.security, config: config.transport)
+        )
+      #endif
+      }
 
-    #if canImport(Network)
-    case .transportServices(let config):
-      let transport = try HTTP2ClientTransport.TransportServices(
-        target: target,
-        transportSecurity: config.security,
-        config: config.transport,
-        serviceConfig: ServiceConfig()
-      )
-      return GRPCClient(transport: NIOClientTransport(transport))
-    #endif
+      try await withGRPCClient(transport: clientTransport) { client in
+        let control = ControlClient(wrapping: client)
+        try await test(control)
+      }
     }
-
   }
 
   private func executeUnaryRPC(control: ControlClient<NIOClientTransport>) async throws {