2 files changed, 36 insertions, 15 deletions

diff --git a/test/cpp/interop/interop_client.cc b/test/cpp/interop/interop_client.cc
index 6340007fa4..05ef0c8d9d 100644
--- a/test/cpp/interop/interop_client.cc
+++ b/test/cpp/interop/interop_client.cc
@@ -80,6 +80,24 @@ CompressionType GetInteropCompressionTypeFromCompressionAlgorithm(
       GPR_ASSERT(false);
   }
 }
+
+void NoopChecks(const InteropClientContextInspector& inspector,
+                const SimpleRequest* request, const SimpleResponse* response) {}
+
+void CompressionChecks(const InteropClientContextInspector& inspector,
+                       const SimpleRequest* request,
+                       const SimpleResponse* response) {
+  GPR_ASSERT(request->response_compression() ==
+             GetInteropCompressionTypeFromCompressionAlgorithm(
+                 inspector.GetCallCompressionAlgorithm()));
+  if (request->response_compression() == NONE) {
+    GPR_ASSERT(!(inspector.GetMessageFlags() & GRPC_WRITE_INTERNAL_COMPRESS));
+  } else if (request->response_type() == PayloadType::COMPRESSABLE) {
+    // requested compression and compressable response => results should always
+    // be compressed.
+    GPR_ASSERT(inspector.GetMessageFlags() & GRPC_WRITE_INTERNAL_COMPRESS);
+  }
+}
 }  // namespace
 
 InteropClient::ServiceStub::ServiceStub(std::shared_ptr<Channel> channel,
@@ -145,9 +163,14 @@ void InteropClient::DoEmpty() {
   gpr_log(GPR_INFO, "Empty rpc done.");
 }
 
-// Shared code to set large payload, make rpc and check response payload.
 void InteropClient::PerformLargeUnary(SimpleRequest* request,
                                       SimpleResponse* response) {
+  PerformLargeUnary(request, response, NoopChecks);
+}
+
+void InteropClient::PerformLargeUnary(SimpleRequest* request,
+                                      SimpleResponse* response,
+                                      CheckerFn custom_checks_fn) {
   ClientContext context;
   InteropClientContextInspector inspector(context);
   // If the request doesn't already specify the response type, default to
@@ -157,21 +180,10 @@ void InteropClient::PerformLargeUnary(SimpleRequest* request,
   request->mutable_payload()->set_body(payload.c_str(), kLargeRequestSize);
 
   Status s = serviceStub_.Get()->UnaryCall(&context, *request, response);
-
-  // Compression related checks.
-  GPR_ASSERT(request->response_compression() ==
-             GetInteropCompressionTypeFromCompressionAlgorithm(
-                 inspector.GetCallCompressionAlgorithm()));
-  if (request->response_compression() == NONE) {
-    GPR_ASSERT(!(inspector.GetMessageFlags() & GRPC_WRITE_INTERNAL_COMPRESS));
-  } else if (request->response_type() == PayloadType::COMPRESSABLE) {
-    // requested compression and compressable response => results should always
-    // be compressed.
-    GPR_ASSERT(inspector.GetMessageFlags() & GRPC_WRITE_INTERNAL_COMPRESS);
-  }
-
   AssertOkOrPrintErrorStatus(s);
 
+  custom_checks_fn(inspector, request, response);
+
   // Payload related checks.
   if (request->response_type() != PayloadType::RANDOM) {
     GPR_ASSERT(response->payload().type() == request->response_type());
@@ -293,7 +305,7 @@ void InteropClient::DoLargeCompressedUnary() {
       SimpleResponse response;
       request.set_response_type(payload_types[i]);
       request.set_response_compression(compression_types[j]);
-      PerformLargeUnary(&request, &response);
+      PerformLargeUnary(&request, &response, CompressionChecks);
       gpr_log(GPR_INFO, "Large compressed unary done %s.", log_suffix);
       gpr_free(log_suffix);
     }
diff --git a/test/cpp/interop/interop_client.h b/test/cpp/interop/interop_client.h
index 1bfb49d514..97a6fd77cf 100644
--- a/test/cpp/interop/interop_client.h
+++ b/test/cpp/interop/interop_client.h
@@ -44,6 +44,11 @@
 namespace grpc {
 namespace testing {
 
+// Function pointer for custom checks.
+using CheckerFn =
+    std::function<void(const InteropClientContextInspector&,
+                       const SimpleRequest*, const SimpleResponse*)>;
+
 class InteropClient {
  public:
   explicit InteropClient(std::shared_ptr<Channel> channel);
@@ -100,6 +105,10 @@ class InteropClient {
   };
 
   void PerformLargeUnary(SimpleRequest* request, SimpleResponse* response);
+
+  /// Run \a custom_check_fn as an additional check.
+  void PerformLargeUnary(SimpleRequest* request, SimpleResponse* response,
+                         CheckerFn custom_checks_fn);
   void AssertOkOrPrintErrorStatus(const Status& s);
   ServiceStub serviceStub_;
 };