From b96c383ef7982f526db190ebba76730fd6c818df Mon Sep 17 00:00:00 2001
From: Penghao Cen <cenph@caicloud.io>
Date: Sat, 13 Jan 2018 05:47:49 +0800
Subject: [PATCH] Check grpc server ready properly

---
 pkg/kubelet/cm/deviceplugin/device_plugin_stub.go |  2 +-
 pkg/kubelet/cm/deviceplugin/endpoint_test.go      | 12 ++----------
 2 files changed, 3 insertions(+), 11 deletions(-)

diff --git a/pkg/kubelet/cm/deviceplugin/device_plugin_stub.go b/pkg/kubelet/cm/deviceplugin/device_plugin_stub.go
index 9969e99989b..5e39dd00853 100644
--- a/pkg/kubelet/cm/deviceplugin/device_plugin_stub.go
+++ b/pkg/kubelet/cm/deviceplugin/device_plugin_stub.go
@@ -70,7 +70,7 @@ func (m *Stub) Start() error {
 	// Wait till grpc server is ready.
 	for i := 0; i < 10; i++ {
 		services := m.server.GetServiceInfo()
-		if len(services) > 1 {
+		if len(services) > 0 {
 			break
 		}
 		time.Sleep(1 * time.Second)
diff --git a/pkg/kubelet/cm/deviceplugin/endpoint_test.go b/pkg/kubelet/cm/deviceplugin/endpoint_test.go
index f4634db85f4..6005310181a 100644
--- a/pkg/kubelet/cm/deviceplugin/endpoint_test.go
+++ b/pkg/kubelet/cm/deviceplugin/endpoint_test.go
@@ -90,20 +90,12 @@ func TestRun(t *testing.T) {
 
 	go e.run()
 	// Wait for the first callback to be issued.
-	select {
-	case <-callbackChan:
-		break
-	}
+	<-callbackChan
 
 	p.Update(updated)
 
 	// Wait for the second callback to be issued.
-	select {
-	case <-callbackChan:
-		break
-	}
-
-	time.Sleep(time.Second)
+	<-callbackChan
 
 	e.mutex.Lock()
 	defer e.mutex.Unlock()