From 847f73b161406d9c01bc27ca178c4299dc826b83 Mon Sep 17 00:00:00 2001
From: Lantao Liu <lantaol@google.com>
Date: Fri, 30 Aug 2019 16:21:15 -0700
Subject: [PATCH] Fix exec test.

Signed-off-by: Lantao Liu <lantaol@google.com>
---
 pkg/validate/streaming.go | 34 +++++++++++++++++++++-------------
 1 file changed, 21 insertions(+), 13 deletions(-)

diff --git a/pkg/validate/streaming.go b/pkg/validate/streaming.go
index d34ceec820..fe4fa8c437 100644
--- a/pkg/validate/streaming.go
+++ b/pkg/validate/streaming.go
@@ -39,6 +39,8 @@ import (
 	. "github.com/onsi/gomega"
 )
 
+const defaultExecStdinCloseTimeout = 20 * time.Second
+
 var _ = framework.KubeDescribe("Streaming", func() {
 	f := framework.NewDefaultCRIFramework()
 
@@ -175,23 +177,29 @@ func createExec(c internalapi.RuntimeService, execReq *runtimeapi.ExecRequest) s
 }
 
 func checkExec(c internalapi.RuntimeService, execServerURL, stdout string, stdoutExactMatch bool, isTty bool) {
-	localOut := &safeBuffer{buffer: bytes.Buffer{}}
-	localErr := &safeBuffer{buffer: bytes.Buffer{}}
-	localInRead, localInWrite := io.Pipe()
-
+	var (
+		localOut                  = &safeBuffer{buffer: bytes.Buffer{}}
+		localErr                  = &safeBuffer{buffer: bytes.Buffer{}}
+		localInRead, localInWrite = io.Pipe()
+		testDone                  = make(chan struct{})
+		wg                        sync.WaitGroup
+	)
+
+	wg.Add(1)
 	// Wait until output read and then shutdown localIn pipe.
 	go func() {
-		ticker := time.NewTicker(5 * time.Second)
-		for {
-			switch {
-			case len(localOut.String()) >= len(stdout):
-				fallthrough
-			case <-ticker.C != time.Time{}:
-				localInWrite.Close()
-				break
-			}
+		defer wg.Done()
+		defer localInWrite.Close()
+		ticker := time.NewTicker(defaultExecStdinCloseTimeout)
+		select {
+		case <-testDone:
+		case <-ticker.C:
 		}
 	}()
+	defer func() {
+		close(testDone)
+		wg.Wait()
+	}()
 
 	// Only http is supported now.
 	// TODO: support streaming APIs via tls.