Review comments

2 files changed, 13 insertions, 9 deletions

diff --git a/test/cpp/end2end/health_service_end2end_test.cc b/test/cpp/end2end/health_service_end2end_test.cc
index 94c92327c8..b96ff53a3e 100644
--- a/test/cpp/end2end/health_service_end2end_test.cc
+++ b/test/cpp/end2end/health_service_end2end_test.cc
@@ -211,14 +211,16 @@ class HealthServiceEnd2endTest : public ::testing::Test {
   // 1. unary client will see NOT_SERVING.
   // 2. unary client still sees NOT_SERVING after a SetServing(true) is called.
   // 3. streaming (Watch) client will see an update.
+  // 4. setting a new service to serving after shutdown will add the service
+  // name but return NOT_SERVING to client.
   // This has to be called last.
   void VerifyHealthCheckServiceShutdown() {
-    const grpc::string kServiceName("service_name");
     HealthCheckServiceInterface* service = server_->GetHealthCheckService();
     EXPECT_TRUE(service != nullptr);
     const grpc::string kHealthyService("healthy_service");
     const grpc::string kUnhealthyService("unhealthy_service");
     const grpc::string kNotRegisteredService("not_registered");
+    const grpc::string kNewService("add_after_shutdown");
     service->SetServingStatus(kHealthyService, true);
     service->SetServingStatus(kUnhealthyService, false);
 
@@ -227,18 +229,12 @@ class HealthServiceEnd2endTest : public ::testing::Test {
     // Start Watch for service.
     ClientContext context;
     HealthCheckRequest request;
-    request.set_service(kServiceName);
+    request.set_service(kHealthyService);
     std::unique_ptr<::grpc::ClientReaderInterface<HealthCheckResponse>> reader =
         hc_stub_->Watch(&context, request);
 
-    // Initial response will be SERVICE_UNKNOWN.
     HealthCheckResponse response;
     EXPECT_TRUE(reader->Read(&response));
-    EXPECT_EQ(response.SERVICE_UNKNOWN, response.status());
-
-    // Set service to SERVING and make sure we get an update.
-    service->SetServingStatus(kServiceName, true);
-    EXPECT_TRUE(reader->Read(&response));
     EXPECT_EQ(response.SERVING, response.status());
 
     SendHealthCheckRpc("", Status::OK, HealthCheckResponse::SERVING);
@@ -248,6 +244,7 @@ class HealthServiceEnd2endTest : public ::testing::Test {
                        HealthCheckResponse::NOT_SERVING);
     SendHealthCheckRpc(kNotRegisteredService,
                        Status(StatusCode::NOT_FOUND, ""));
+    SendHealthCheckRpc(kNewService, Status(StatusCode::NOT_FOUND, ""));
 
     // Shutdown health check service.
     service->Shutdown();
@@ -270,6 +267,12 @@ class HealthServiceEnd2endTest : public ::testing::Test {
     service->SetServingStatus(kHealthyService, true);
     SendHealthCheckRpc(kHealthyService, Status::OK,
                        HealthCheckResponse::NOT_SERVING);
+
+    // Adding serving status for a new service after shutdown will return
+    // NOT_SERVING.
+    service->SetServingStatus(kNewService, true);
+    SendHealthCheckRpc(kNewService, Status::OK,
+                       HealthCheckResponse::NOT_SERVING);
   }
 
   TestServiceImpl echo_test_service_;
diff --git a/test/cpp/end2end/test_health_check_service_impl.cc b/test/cpp/end2end/test_health_check_service_impl.cc
index fa70a44d24..0801e30199 100644
--- a/test/cpp/end2end/test_health_check_service_impl.cc
+++ b/test/cpp/end2end/test_health_check_service_impl.cc
@@ -37,6 +37,7 @@ Status HealthCheckServiceImpl::Check(ServerContext* context,
   response->set_status(iter->second);
   return Status::OK;
 }
+
 Status HealthCheckServiceImpl::Watch(
     ServerContext* context, const HealthCheckRequest* request,
     ::grpc::ServerWriter<HealthCheckResponse>* writer) {
@@ -66,7 +67,7 @@ void HealthCheckServiceImpl::SetStatus(
     HealthCheckResponse::ServingStatus status) {
   std::lock_guard<std::mutex> lock(mu_);
   if (shutdown_) {
-    return;
+    status = HealthCheckResponse::NOT_SERVING;
   }
   status_map_[service_name] = status;
 }