Added support for Error Responses from the Reflection Service (#1682)

Motivation:

Whenever an error occurs during the exchange between the server implementing the Reflection Service and a client, the client should receive an Error Message.

Modifications:

Replaced throwing errors with creating an Error Reponse and sending it back to the client. Also, changed the APIs to use the Result type and added new extensions for the Reflection_ServerReflectionResponse initialisation. Also modified the tests.

Result:

When an error occurs within the Reflection Service, the clients will be sent an Error Response.

Sources/GRPCReflectionService/Server/ReflectionService.swift

@@ -129,7 +129,7 @@ internal struct ReflectionServiceData: Sendable {
 
   internal func serialisedFileDescriptorProtosForDependenciesOfFile(
     named fileName: String
-  ) throws -> [Data] {
+  ) -> Result<[Data], GRPCStatus> {
     var toVisit = Deque<String>()
     var visited = Set<String>()
     var serializedFileDescriptorProtos: [Data] = []
@@ -147,37 +147,59 @@ internal struct ReflectionServiceData: Sendable {
         let serializedFileDescriptorProto = protoData.serializedFileDescriptorProto
         serializedFileDescriptorProtos.append(serializedFileDescriptorProto)
       } else {
-        throw GRPCStatus(
-          code: .notFound,
-          message: "The provided file or a dependency of the provided file could not be found."
+        return .failure(
+          GRPCStatus(
+            code: .notFound,
+            message: "The provided file or a dependency of the provided file could not be found."
+          )
         )
       }
       visited.insert(currentFileName)
     }
-    return serializedFileDescriptorProtos
+    return .success(serializedFileDescriptorProtos)
   }
 
-  internal func nameOfFileContainingSymbol(named symbolName: String) -> String? {
-    return self.fileNameBySymbol[symbolName]
+  internal func nameOfFileContainingSymbol(named symbolName: String) -> Result<String, GRPCStatus> {
+    guard let fileName = self.fileNameBySymbol[symbolName] else {
+      return .failure(
+        GRPCStatus(
+          code: .notFound,
+          message: "The provided symbol could not be found."
+        )
+      )
+    }
+    return .success(fileName)
   }
 
   internal func nameOfFileContainingExtension(
     extendeeName: String,
     fieldNumber number: Int32
-  ) -> String? {
+  ) -> Result<String, GRPCStatus> {
     let key = ExtensionDescriptor(extendeeTypeName: extendeeName, fieldNumber: number)
-    return self.fileNameByExtensionDescriptor[key]
+    guard let fileName = self.fileNameByExtensionDescriptor[key] else {
+      return .failure(
+        GRPCStatus(
+          code: .notFound,
+          message: "The provided extension could not be found."
+        )
+      )
+    }
+    return .success(fileName)
   }
 
   // Returns an empty array if the type has no extensions.
-  internal func extensionsFieldNumbersOfType(named typeName: String) throws -> [Int32] {
+  internal func extensionsFieldNumbersOfType(
+    named typeName: String
+  ) -> Result<[Int32], GRPCStatus> {
     guard let fieldNumbers = self.fieldNumbersByType[typeName] else {
-      throw GRPCStatus(
-        code: .invalidArgument,
-        message: "The provided type is invalid."
+      return .failure(
+        GRPCStatus(
+          code: .invalidArgument,
+          message: "The provided type is invalid."
+        )
       )
     }
-    return fieldNumbers
+    return .success(fieldNumbers)
   }
 }
 
@@ -191,17 +213,26 @@ internal final class ReflectionServiceProvider: Reflection_ServerReflectionAsync
     )
   }
 
+  internal func _findFileByFileName(
+    _ fileName: String
+  ) -> Result<Reflection_ServerReflectionResponse.OneOf_MessageResponse, GRPCStatus> {
+    return self.protoRegistry
+      .serialisedFileDescriptorProtosForDependenciesOfFile(named: fileName)
+      .map { fileDescriptorProtos in
+        Reflection_ServerReflectionResponse.OneOf_MessageResponse.fileDescriptorResponse(
+          .with {
+            $0.fileDescriptorProto = fileDescriptorProtos
+          }
+        )
+      }
+  }
+
   internal func findFileByFileName(
     _ fileName: String,
     request: Reflection_ServerReflectionRequest
-  ) throws -> Reflection_ServerReflectionResponse {
-    return Reflection_ServerReflectionResponse(
-      request: request,
-      fileDescriptorResponse: try .with {
-        $0.fileDescriptorProto = try self.protoRegistry
-          .serialisedFileDescriptorProtosForDependenciesOfFile(named: fileName)
-      }
-    )
+  ) -> Reflection_ServerReflectionResponse {
+    let result = self._findFileByFileName(fileName)
+    return result.makeResponse(request: request)
   }
 
   internal func getServicesNames(
@@ -215,53 +246,50 @@ internal final class ReflectionServiceProvider: Reflection_ServerReflectionAsync
     }
     return Reflection_ServerReflectionResponse(
       request: request,
-      listServicesResponse: listServicesResponse
+      messageResponse: .listServicesResponse(listServicesResponse)
     )
   }
 
   internal func findFileBySymbol(
     _ symbolName: String,
     request: Reflection_ServerReflectionRequest
-  ) throws -> Reflection_ServerReflectionResponse {
-    guard let fileName = self.protoRegistry.nameOfFileContainingSymbol(named: symbolName) else {
-      throw GRPCStatus(
-        code: .notFound,
-        message: "The provided symbol could not be found."
-      )
+  ) -> Reflection_ServerReflectionResponse {
+    let result = self.protoRegistry.nameOfFileContainingSymbol(
+      named: symbolName
+    ).flatMap {
+      self._findFileByFileName($0)
     }
-    return try self.findFileByFileName(fileName, request: request)
+    return result.makeResponse(request: request)
   }
 
   internal func findFileByExtension(
     extensionRequest: Reflection_ExtensionRequest,
     request: Reflection_ServerReflectionRequest
-  ) throws -> Reflection_ServerReflectionResponse {
-    guard
-      let fileName = self.protoRegistry.nameOfFileContainingExtension(
-        extendeeName: extensionRequest.containingType,
-        fieldNumber: extensionRequest.extensionNumber
-      )
-    else {
-      throw GRPCStatus(
-        code: .notFound,
-        message: "The provided extension could not be found."
-      )
+  ) -> Reflection_ServerReflectionResponse {
+    let result = self.protoRegistry.nameOfFileContainingExtension(
+      extendeeName: extensionRequest.containingType,
+      fieldNumber: extensionRequest.extensionNumber
+    ).flatMap {
+      self._findFileByFileName($0)
     }
-    return try self.findFileByFileName(fileName, request: request)
+    return result.makeResponse(request: request)
   }
 
   internal func findExtensionsFieldNumbersOfType(
     named typeName: String,
     request: Reflection_ServerReflectionRequest
-  ) throws -> Reflection_ServerReflectionResponse {
-    let fieldNumbers = try self.protoRegistry.extensionsFieldNumbersOfType(named: typeName)
-    return Reflection_ServerReflectionResponse(
-      request: request,
-      extensionNumberResponse: .with {
-        $0.baseTypeName = typeName
-        $0.extensionNumber = fieldNumbers
-      }
-    )
+  ) -> Reflection_ServerReflectionResponse {
+    let result = self.protoRegistry.extensionsFieldNumbersOfType(
+      named: typeName
+    ).map { fieldNumbers in
+      Reflection_ServerReflectionResponse.OneOf_MessageResponse.allExtensionNumbersResponse(
+        Reflection_ExtensionNumberResponse.with {
+          $0.baseTypeName = typeName
+          $0.extensionNumber = fieldNumbers
+        }
+      )
+    }
+    return result.makeResponse(request: request)
   }
 
   internal func serverReflectionInfo(
@@ -272,7 +300,7 @@ internal final class ReflectionServiceProvider: Reflection_ServerReflectionAsync
     for try await request in requestStream {
       switch request.messageRequest {
       case let .fileByFilename(fileName):
-        let response = try self.findFileByFileName(
+        let response = self.findFileByFileName(
           fileName,
           request: request
         )
@@ -283,28 +311,37 @@ internal final class ReflectionServiceProvider: Reflection_ServerReflectionAsync
         try await responseStream.send(response)
 
       case let .fileContainingSymbol(symbolName):
-        let response = try self.findFileBySymbol(
+        let response = self.findFileBySymbol(
           symbolName,
           request: request
         )
         try await responseStream.send(response)
 
       case let .fileContainingExtension(extensionRequest):
-        let response = try self.findFileByExtension(
+        let response = self.findFileByExtension(
           extensionRequest: extensionRequest,
           request: request
         )
         try await responseStream.send(response)
 
       case let .allExtensionNumbersOfType(typeName):
-        let response = try self.findExtensionsFieldNumbersOfType(
+        let response = self.findExtensionsFieldNumbersOfType(
           named: typeName,
           request: request
         )
         try await responseStream.send(response)
 
       default:
-        throw GRPCStatus(code: .unimplemented)
+        let response = Reflection_ServerReflectionResponse(
+          request: request,
+          messageResponse: .errorResponse(
+            Reflection_ErrorResponse.with {
+              $0.errorCode = Int32(GRPCStatus.Code.unimplemented.rawValue)
+              $0.errorMessage = "The request is not implemented."
+            }
+          )
+        )
+        try await responseStream.send(response)
       }
     }
   }
@@ -313,34 +350,12 @@ internal final class ReflectionServiceProvider: Reflection_ServerReflectionAsync
 extension Reflection_ServerReflectionResponse {
   init(
     request: Reflection_ServerReflectionRequest,
-    fileDescriptorResponse: Reflection_FileDescriptorResponse
-  ) {
-    self = .with {
-      $0.validHost = request.host
-      $0.originalRequest = request
-      $0.fileDescriptorResponse = fileDescriptorResponse
-    }
-  }
-
-  init(
-    request: Reflection_ServerReflectionRequest,
-    listServicesResponse: Reflection_ListServiceResponse
-  ) {
-    self = .with {
-      $0.validHost = request.host
-      $0.originalRequest = request
-      $0.listServicesResponse = listServicesResponse
-    }
-  }
-
-  init(
-    request: Reflection_ServerReflectionRequest,
-    extensionNumberResponse: Reflection_ExtensionNumberResponse
+    messageResponse: Reflection_ServerReflectionResponse.OneOf_MessageResponse
   ) {
     self = .with {
       $0.validHost = request.host
       $0.originalRequest = request
-      $0.allExtensionNumbersResponse = extensionNumberResponse
+      $0.messageResponse = messageResponse
     }
   }
 }
@@ -378,3 +393,33 @@ extension Google_Protobuf_FileDescriptorProto {
     return names
   }
 }
+
+extension Result<Reflection_ServerReflectionResponse.OneOf_MessageResponse, GRPCStatus> {
+  func recover() -> Result<Reflection_ServerReflectionResponse.OneOf_MessageResponse, Never> {
+    self.flatMapError { status in
+      let error = Reflection_ErrorResponse.with {
+        $0.errorCode = Int32(status.code.rawValue)
+        $0.errorMessage = status.message ?? ""
+      }
+      return .success(.errorResponse(error))
+    }
+  }
+
+  func makeResponse(
+    request: Reflection_ServerReflectionRequest
+  ) -> Reflection_ServerReflectionResponse {
+    let result = self.recover().attachRequest(request)
+    // Safe to '!' as the failure type is 'Never'.
+    return try! result.get()
+  }
+}
+
+extension Result where Success == Reflection_ServerReflectionResponse.OneOf_MessageResponse {
+  func attachRequest(
+    _ request: Reflection_ServerReflectionRequest
+  ) -> Result<Reflection_ServerReflectionResponse, Failure> {
+    self.map { message in
+      Reflection_ServerReflectionResponse(request: request, messageResponse: message)
+    }
+  }
+}

Tests/GRPCTests/GRPCReflectionServiceTests/ReflectionServiceIntegrationTests.swift

@@ -259,4 +259,90 @@ final class ReflectionServiceIntegrationTests: GRPCTestCase {
     XCTAssertEqual(message.allExtensionNumbersResponse.baseTypeName, "packagebar2.inputMessage2")
     XCTAssertEqual(message.allExtensionNumbersResponse.extensionNumber, [1, 2, 3, 4, 5])
   }
+
+  func testErrorResponseFileByFileNameRequest() async throws {
+    try self.setUpServerAndChannel()
+    let client = Reflection_ServerReflectionAsyncClient(channel: self.channel!)
+    let serviceReflectionInfo = client.makeServerReflectionInfoCall()
+    try await serviceReflectionInfo.requestStream.send(
+      .with {
+        $0.host = "127.0.0.1"
+        $0.fileByFilename = "invalidFileName.proto"
+      }
+    )
+    serviceReflectionInfo.requestStream.finish()
+    var iterator = serviceReflectionInfo.responseStream.makeAsyncIterator()
+    guard let message = try await iterator.next() else {
+      return XCTFail("Could not get a response message.")
+    }
+
+    XCTAssertEqual(message.errorResponse.errorCode, Int32(GRPCStatus.Code.notFound.rawValue))
+    XCTAssertEqual(
+      message.errorResponse.errorMessage,
+      "The provided file or a dependency of the provided file could not be found."
+    )
+  }
+
+  func testErrorResponseFileBySymbolRequest() async throws {
+    try self.setUpServerAndChannel()
+    let client = Reflection_ServerReflectionAsyncClient(channel: self.channel!)
+    let serviceReflectionInfo = client.makeServerReflectionInfoCall()
+    try await serviceReflectionInfo.requestStream.send(
+      .with {
+        $0.host = "127.0.0.1"
+        $0.fileContainingSymbol = "packagebar1.invalidEnumType1"
+      }
+    )
+    serviceReflectionInfo.requestStream.finish()
+    var iterator = serviceReflectionInfo.responseStream.makeAsyncIterator()
+    guard let message = try await iterator.next() else {
+      return XCTFail("Could not get a response message.")
+    }
+
+    XCTAssertEqual(message.errorResponse.errorCode, Int32(GRPCStatus.Code.notFound.rawValue))
+    XCTAssertEqual(message.errorResponse.errorMessage, "The provided symbol could not be found.")
+  }
+
+  func testErrorResponseFileByExtensionRequest() async throws {
+    try self.setUpServerAndChannel()
+    let client = Reflection_ServerReflectionAsyncClient(channel: self.channel!)
+    let serviceReflectionInfo = client.makeServerReflectionInfoCall()
+    try await serviceReflectionInfo.requestStream.send(
+      .with {
+        $0.host = "127.0.0.1"
+        $0.fileContainingExtension = .with {
+          $0.containingType = "packagebar1.invalidInputMessage1"
+          $0.extensionNumber = 2
+        }
+      }
+    )
+    serviceReflectionInfo.requestStream.finish()
+    var iterator = serviceReflectionInfo.responseStream.makeAsyncIterator()
+    guard let message = try await iterator.next() else {
+      return XCTFail("Could not get a response message.")
+    }
+
+    XCTAssertEqual(message.errorResponse.errorCode, Int32(GRPCStatus.Code.notFound.rawValue))
+    XCTAssertEqual(message.errorResponse.errorMessage, "The provided extension could not be found.")
+  }
+
+  func testErrorResponseAllExtensionNumbersOfTypeRequest() async throws {
+    try self.setUpServerAndChannel()
+    let client = Reflection_ServerReflectionAsyncClient(channel: self.channel!)
+    let serviceReflectionInfo = client.makeServerReflectionInfoCall()
+    try await serviceReflectionInfo.requestStream.send(
+      .with {
+        $0.host = "127.0.0.1"
+        $0.allExtensionNumbersOfType = "packagebar2.invalidInputMessage2"
+      }
+    )
+    serviceReflectionInfo.requestStream.finish()
+    var iterator = serviceReflectionInfo.responseStream.makeAsyncIterator()
+    guard let message = try await iterator.next() else {
+      return XCTFail("Could not get a response message.")
+    }
+
+    XCTAssertEqual(message.errorResponse.errorCode, Int32(GRPCStatus.Code.invalidArgument.rawValue))
+    XCTAssertEqual(message.errorResponse.errorMessage, "The provided type is invalid.")
+  }
 }

Tests/GRPCTests/GRPCReflectionServiceTests/ReflectionServiceUnitTests.swift

@@ -117,38 +117,52 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
   func testNameOfFileContainingSymbolEnum() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let fileName = registry.nameOfFileContainingSymbol(named: "packagebar2.enumType2")
-    XCTAssertEqual(fileName, "bar2.proto")
+    let nameOfFileContainingSymbolResult = registry.nameOfFileContainingSymbol(
+      named: "packagebar2.enumType2"
+    )
+    XCTAssertEqual(try nameOfFileContainingSymbolResult.get(), "bar2.proto")
   }
 
   func testNameOfFileContainingSymbolMessage() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let fileName = registry.nameOfFileContainingSymbol(named: "packagebar1.inputMessage1")
-    XCTAssertEqual(fileName, "bar1.proto")
+    let nameOfFileContainingSymbolResult = registry.nameOfFileContainingSymbol(
+      named: "packagebar1.inputMessage1"
+    )
+    XCTAssertEqual(try nameOfFileContainingSymbolResult.get(), "bar1.proto")
   }
 
   func testNameOfFileContainingSymbolService() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let fileName = registry.nameOfFileContainingSymbol(named: "packagebar3.service3")
-    XCTAssertEqual(fileName, "bar3.proto")
+    let nameOfFileContainingSymbolResult = registry.nameOfFileContainingSymbol(
+      named: "packagebar3.service3"
+    )
+    XCTAssertEqual(try nameOfFileContainingSymbolResult.get(), "bar3.proto")
   }
 
   func testNameOfFileContainingSymbolMethod() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let fileName = registry.nameOfFileContainingSymbol(
+    let nameOfFileContainingSymbolResult = registry.nameOfFileContainingSymbol(
       named: "packagebar4.service4.testMethod4"
     )
-    XCTAssertEqual(fileName, "bar4.proto")
+    XCTAssertEqual(try nameOfFileContainingSymbolResult.get(), "bar4.proto")
   }
 
   func testNameOfFileContainingSymbolNonExistentSymbol() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let fileName = registry.nameOfFileContainingSymbol(named: "packagebar2.enumType3")
-    XCTAssertNil(fileName)
+    let nameOfFileContainingSymbolResult = registry.nameOfFileContainingSymbol(
+      named: "packagebar2.enumType3"
+    )
+    XCTAssertThrowsGRPCStatus(try nameOfFileContainingSymbolResult.get()) {
+      status in
+      XCTAssertEqual(
+        status,
+        GRPCStatus(code: .notFound, message: "The provided symbol could not be found.")
+      )
+    }
   }
 
   // Testing the serializedFileDescriptorProto method in different cases.
@@ -156,95 +170,112 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
   func testSerialisedFileDescriptorProtosForDependenciesOfFile() throws {
     var protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let serializedFileDescriptorProtos =
-      try registry
+    let serializedFileDescriptorProtosResult =
+      registry
       .serialisedFileDescriptorProtosForDependenciesOfFile(named: "bar1.proto")
-    let fileDescriptorProtos = try serializedFileDescriptorProtos.map {
-      try Google_Protobuf_FileDescriptorProto(serializedData: $0)
-    }
-    // Tests that the functions returns all the transitive dependencies, with their services and
-    // methods, together with the initial proto, as serialized data.
-    XCTAssertEqual(fileDescriptorProtos.count, 4)
-    for fileDescriptorProto in fileDescriptorProtos {
-      guard let protoIndex = protos.firstIndex(of: fileDescriptorProto) else {
-        return XCTFail(
-          """
-          Could not find the file descriptor proto of \(fileDescriptorProto.name) \
-          in the original file descriptor protos list.
-          """
-        )
-      }
 
-      for service in fileDescriptorProto.service {
-        guard let serviceIndex = protos[protoIndex].service.firstIndex(of: service) else {
+    switch serializedFileDescriptorProtosResult {
+    case .success(let serializedFileDescriptorProtos):
+      let fileDescriptorProtos = try serializedFileDescriptorProtos.map {
+        try Google_Protobuf_FileDescriptorProto(serializedData: $0)
+      }
+      // Tests that the functions returns all the transitive dependencies, with their services and
+      // methods, together with the initial proto, as serialized data.
+      XCTAssertEqual(fileDescriptorProtos.count, 4)
+      for fileDescriptorProto in fileDescriptorProtos {
+        guard let protoIndex = protos.firstIndex(of: fileDescriptorProto) else {
           return XCTFail(
             """
-            Could not find the \(service.name) in the service \
-            list of the \(fileDescriptorProto.name) file descriptor proto.
+            Could not find the file descriptor proto of \(fileDescriptorProto.name) \
+            in the original file descriptor protos list.
             """
           )
         }
 
-        let originalMethods = protos[protoIndex].service[serviceIndex].method
-        for method in service.method {
-          XCTAssert(originalMethods.contains(method))
-        }
+        for service in fileDescriptorProto.service {
+          guard let serviceIndex = protos[protoIndex].service.firstIndex(of: service) else {
+            return XCTFail(
+              """
+              Could not find the \(service.name) in the service \
+              list of the \(fileDescriptorProto.name) file descriptor proto.
+              """
+            )
+          }
 
-        for messageType in fileDescriptorProto.messageType {
-          XCTAssert(protos[protoIndex].messageType.contains(messageType))
+          let originalMethods = protos[protoIndex].service[serviceIndex].method
+          for method in service.method {
+            XCTAssert(originalMethods.contains(method))
+          }
+
+          for messageType in fileDescriptorProto.messageType {
+            XCTAssert(protos[protoIndex].messageType.contains(messageType))
+          }
         }
-      }
 
-      protos.removeAll { $0 == fileDescriptorProto }
+        protos.removeAll { $0 == fileDescriptorProto }
+      }
+      XCTAssert(protos.isEmpty)
+    case .failure(let status):
+      XCTFail(
+        "Faild with GRPCStatus code: " + String(status.code.rawValue) + " and message: "
+          + (status.message ?? "empty") + "."
+      )
     }
-    XCTAssert(protos.isEmpty)
   }
 
   func testSerialisedFileDescriptorProtosForDependenciesOfFileComplexDependencyGraph() throws {
     var protos = makeProtosWithComplexDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let serializedFileDescriptorProtos =
-      try registry
+    let serializedFileDescriptorProtosResult =
+      registry
       .serialisedFileDescriptorProtosForDependenciesOfFile(named: "foo0.proto")
-    let fileDescriptorProtos = try serializedFileDescriptorProtos.map {
-      try Google_Protobuf_FileDescriptorProto(serializedData: $0)
-    }
-    // Tests that the functions returns all the tranzitive dependencies, with their services and
-    // methods, together with the initial proto, as serialized data.
-    XCTAssertEqual(fileDescriptorProtos.count, 21)
-    for fileDescriptorProto in fileDescriptorProtos {
-      guard let protoIndex = protos.firstIndex(of: fileDescriptorProto) else {
-        return XCTFail(
-          """
-          Could not find the file descriptor proto of \(fileDescriptorProto.name) \
-          in the original file descriptor protos list.
-          """
-        )
+    switch serializedFileDescriptorProtosResult {
+    case .success(let serializedFileDescriptorProtos):
+      let fileDescriptorProtos = try serializedFileDescriptorProtos.map {
+        try Google_Protobuf_FileDescriptorProto(serializedData: $0)
       }
-
-      for service in fileDescriptorProto.service {
-        guard let serviceIndex = protos[protoIndex].service.firstIndex(of: service) else {
+      // Tests that the functions returns all the tranzitive dependencies, with their services and
+      // methods, together with the initial proto, as serialized data.
+      XCTAssertEqual(fileDescriptorProtos.count, 21)
+      for fileDescriptorProto in fileDescriptorProtos {
+        guard let protoIndex = protos.firstIndex(of: fileDescriptorProto) else {
           return XCTFail(
             """
-            Could not find the \(service.name) in the service \
-            list of the \(fileDescriptorProto.name) file descriptor proto.
+            Could not find the file descriptor proto of \(fileDescriptorProto.name) \
+            in the original file descriptor protos list.
             """
           )
         }
 
-        let originalMethods = protos[protoIndex].service[serviceIndex].method
-        for method in service.method {
-          XCTAssert(originalMethods.contains(method))
-        }
+        for service in fileDescriptorProto.service {
+          guard let serviceIndex = protos[protoIndex].service.firstIndex(of: service) else {
+            return XCTFail(
+              """
+              Could not find the \(service.name) in the service \
+              list of the \(fileDescriptorProto.name) file descriptor proto.
+              """
+            )
+          }
 
-        for messageType in fileDescriptorProto.messageType {
-          XCTAssert(protos[protoIndex].messageType.contains(messageType))
+          let originalMethods = protos[protoIndex].service[serviceIndex].method
+          for method in service.method {
+            XCTAssert(originalMethods.contains(method))
+          }
+
+          for messageType in fileDescriptorProto.messageType {
+            XCTAssert(protos[protoIndex].messageType.contains(messageType))
+          }
         }
-      }
 
-      protos.removeAll { $0 == fileDescriptorProto }
+        protos.removeAll { $0 == fileDescriptorProto }
+      }
+      XCTAssert(protos.isEmpty)
+    case .failure(let status):
+      XCTFail(
+        "Faild with GRPCStatus code: " + String(status.code.rawValue) + " and message: "
+          + (status.message ?? "empty") + "."
+      )
     }
-    XCTAssert(protos.isEmpty)
   }
 
   func testSerialisedFileDescriptorProtosForDependenciesOfFileDependencyLoops() throws {
@@ -254,57 +285,67 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     protos[2].dependency.append("bar1.proto")
     protos[3].dependency.append("bar1.proto")
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let serializedFileDescriptorProtos =
-      try registry
+    let serializedFileDescriptorProtosResult =
+      registry
       .serialisedFileDescriptorProtosForDependenciesOfFile(named: "bar1.proto")
-    let fileDescriptorProtos = try serializedFileDescriptorProtos.map {
-      try Google_Protobuf_FileDescriptorProto(serializedData: $0)
-    }
-    // Test that we get only 4 serialized File Descriptor Protos as response.
-    XCTAssertEqual(fileDescriptorProtos.count, 4)
-    for fileDescriptorProto in fileDescriptorProtos {
-      guard let protoIndex = protos.firstIndex(of: fileDescriptorProto) else {
-        return XCTFail(
-          """
-          Could not find the file descriptor proto of \(fileDescriptorProto.name) \
-          in the original file descriptor protos list.
-          """
-        )
+    switch serializedFileDescriptorProtosResult {
+    case .success(let serializedFileDescriptorProtos):
+      let fileDescriptorProtos = try serializedFileDescriptorProtos.map {
+        try Google_Protobuf_FileDescriptorProto(serializedData: $0)
       }
-
-      for service in fileDescriptorProto.service {
-        guard let serviceIndex = protos[protoIndex].service.firstIndex(of: service) else {
+      // Test that we get only 4 serialized File Descriptor Protos as response.
+      XCTAssertEqual(fileDescriptorProtos.count, 4)
+      for fileDescriptorProto in fileDescriptorProtos {
+        guard let protoIndex = protos.firstIndex(of: fileDescriptorProto) else {
           return XCTFail(
             """
-            Could not find the \(service.name) in the service \
-            list of the \(fileDescriptorProto.name) file descriptor proto.
+            Could not find the file descriptor proto of \(fileDescriptorProto.name) \
+            in the original file descriptor protos list.
             """
           )
         }
 
-        let originalMethods = protos[protoIndex].service[serviceIndex].method
-        for method in service.method {
-          XCTAssert(originalMethods.contains(method))
-        }
+        for service in fileDescriptorProto.service {
+          guard let serviceIndex = protos[protoIndex].service.firstIndex(of: service) else {
+            return XCTFail(
+              """
+              Could not find the \(service.name) in the service \
+              list of the \(fileDescriptorProto.name) file descriptor proto.
+              """
+            )
+          }
 
-        for messageType in fileDescriptorProto.messageType {
-          XCTAssert(protos[protoIndex].messageType.contains(messageType))
+          let originalMethods = protos[protoIndex].service[serviceIndex].method
+          for method in service.method {
+            XCTAssert(originalMethods.contains(method))
+          }
+
+          for messageType in fileDescriptorProto.messageType {
+            XCTAssert(protos[protoIndex].messageType.contains(messageType))
+          }
         }
-      }
 
-      protos.removeAll { $0 == fileDescriptorProto }
+        protos.removeAll { $0 == fileDescriptorProto }
+      }
+      XCTAssert(protos.isEmpty)
+    case .failure(let status):
+      XCTFail(
+        "Faild with GRPCStatus code: " + String(status.code.rawValue) + " and message: "
+          + (status.message ?? "empty") + "."
+      )
     }
-    XCTAssert(protos.isEmpty)
   }
 
   func testSerialisedFileDescriptorProtosForDependenciesOfFileInvalidFile() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    XCTAssertThrowsError(
-      try registry.serialisedFileDescriptorProtosForDependenciesOfFile(named: "invalid.proto")
-    ) { error in
+    let serializedFileDescriptorProtosForDependenciesOfFileResult =
+      registry.serialisedFileDescriptorProtosForDependenciesOfFile(named: "invalid.proto")
+
+    XCTAssertThrowsGRPCStatus(try serializedFileDescriptorProtosForDependenciesOfFileResult.get()) {
+      status in
       XCTAssertEqual(
-        error as? GRPCStatus,
+        status,
         GRPCStatus(
           code: .notFound,
           message: "The provided file or a dependency of the provided file could not be found."
@@ -317,11 +358,13 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     var protos = makeProtosWithDependencies()
     protos[0].dependency.append("invalidDependency")
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    XCTAssertThrowsError(
-      try registry.serialisedFileDescriptorProtosForDependenciesOfFile(named: "bar1.proto")
-    ) { error in
+    let serializedFileDescriptorProtosForDependenciesOfFileResult =
+      registry.serialisedFileDescriptorProtosForDependenciesOfFile(named: "bar1.proto")
+
+    XCTAssertThrowsGRPCStatus(try serializedFileDescriptorProtosForDependenciesOfFileResult.get()) {
+      status in
       XCTAssertEqual(
-        error as? GRPCStatus,
+        status,
         GRPCStatus(
           code: .notFound,
           message: "The provided file or a dependency of the provided file could not be found."
@@ -338,11 +381,11 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     for proto in protos {
       for `extension` in proto.extension {
         let typeName = String(`extension`.extendee.drop(while: { $0 == "." }))
-        let registryFileName = registry.nameOfFileContainingExtension(
+        let registryFileNameResult = registry.nameOfFileContainingExtension(
           extendeeName: typeName,
           fieldNumber: `extension`.number
         )
-        XCTAssertEqual(registryFileName, proto.name)
+        XCTAssertEqual(try registryFileNameResult.get(), proto.name)
       }
     }
   }
@@ -350,21 +393,35 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
   func testNameOfFileContainingExtensionsInvalidTypeName() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let registryFileName = registry.nameOfFileContainingExtension(
+    let registryFileNameResult = registry.nameOfFileContainingExtension(
       extendeeName: "InvalidType",
       fieldNumber: 2
     )
-    XCTAssertNil(registryFileName)
+
+    XCTAssertThrowsGRPCStatus(try registryFileNameResult.get()) {
+      status in
+      XCTAssertEqual(
+        status,
+        GRPCStatus(code: .notFound, message: "The provided extension could not be found.")
+      )
+    }
   }
 
   func testNameOfFileContainingExtensionsInvalidFieldNumber() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let registryFileName = registry.nameOfFileContainingExtension(
+    let registryFileNameResult = registry.nameOfFileContainingExtension(
       extendeeName: protos[0].extension[0].extendee,
       fieldNumber: 9
     )
-    XCTAssertNil(registryFileName)
+
+    XCTAssertThrowsGRPCStatus(try registryFileNameResult.get()) {
+      status in
+      XCTAssertEqual(
+        status,
+        GRPCStatus(code: .notFound, message: "The provided extension could not be found.")
+      )
+    }
   }
 
   func testNameOfFileContainingExtensionsDuplicatedExtensions() throws {
@@ -403,10 +460,11 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
       }
     )
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let extensionNumbers = try registry.extensionsFieldNumbersOfType(
+    let extensionsFieldNumbersOfTypeResult = registry.extensionsFieldNumbersOfType(
       named: "packagebar1.inputMessage1"
     )
-    XCTAssertEqual(extensionNumbers, [1, 2, 3, 4, 5, 120])
+
+    XCTAssertEqual(try extensionsFieldNumbersOfTypeResult.get(), [1, 2, 3, 4, 5, 120])
   }
 
   func testExtensionsFieldNumbersOfTypeNoExtensionsType() throws {
@@ -423,26 +481,25 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
       }
     )
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let extensionNumbers = try registry.extensionsFieldNumbersOfType(
+    let extensionsFieldNumbersOfTypeResult = registry.extensionsFieldNumbersOfType(
       named: "packagebar1.noExtensionMessage"
     )
-    XCTAssertEqual(extensionNumbers, [])
+
+    XCTAssertEqual(try extensionsFieldNumbersOfTypeResult.get(), [])
   }
 
   func testExtensionsFieldNumbersOfTypeInvalidTypeName() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    XCTAssertThrowsError(
-      try registry.extensionsFieldNumbersOfType(
-        named: "packagebar1.invalidTypeMessage"
-      )
-    ) { error in
+    let extensionsFieldNumbersOfTypeResult = registry.extensionsFieldNumbersOfType(
+      named: "packagebar1.invalidTypeMessage"
+    )
+
+    XCTAssertThrowsGRPCStatus(try extensionsFieldNumbersOfTypeResult.get()) {
+      status in
       XCTAssertEqual(
-        error as? GRPCStatus,
-        GRPCStatus(
-          code: .invalidArgument,
-          message: "The provided type is invalid."
-        )
+        status,
+        GRPCStatus(code: .invalidArgument, message: "The provided type is invalid.")
       )
     }
   }
@@ -456,9 +513,10 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
       }
     )
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let extensionNumbers = try registry.extensionsFieldNumbersOfType(
+    let extensionsFieldNumbersOfTypeResult = registry.extensionsFieldNumbersOfType(
       named: "packagebar1.inputMessage1"
     )
-    XCTAssertEqual(extensionNumbers, [1, 2, 3, 4, 5, 130])
+
+    XCTAssertEqual(try extensionsFieldNumbersOfTypeResult.get(), [1, 2, 3, 4, 5, 130])
   }
 }

Tests/GRPCTests/GRPCReflectionServiceTests/Utils.swift

@@ -17,6 +17,7 @@
 import Foundation
 import GRPC
 import SwiftProtobuf
+import XCTest
 
 internal func makeExtensions(
   forType typeName: String,
@@ -136,6 +137,19 @@ internal func makeProtosWithComplexDependencies() -> [Google_Protobuf_FileDescri
   return protos
 }
 
+func XCTAssertThrowsGRPCStatus<T>(
+  _ expression: @autoclosure () throws -> T,
+  _ errorHandler: (GRPCStatus) -> Void
+) {
+  XCTAssertThrowsError(try expression()) { error in
+    guard let error = error as? GRPCStatus else {
+      return XCTFail("Error had unexpected type '\(type(of: error))'")
+    }
+
+    errorHandler(error)
+  }
+}
+
 extension Sequence where Element == Google_Protobuf_FileDescriptorProto {
   var serviceNames: [String] {
     self.flatMap { $0.service.map { $0.name } }