Deflake remote command execution test

pkg/client/remotecommand/remotecommand_test.go

@@ -18,271 +18,175 @@ package remotecommand
 
 import (
 	"bytes"
-	"errors"
+	"fmt"
-	"io"
+	"io/ioutil"
 	"net/http"
+	"net/http/httptest"
+	"net/url"
 	"strings"
-	"sync"
 	"testing"
-	"time"
 
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/util/httpstream"
+	"github.com/GoogleCloudPlatform/kubernetes/pkg/util/httpstream/spdy"
 )
 
-type fakeUpgrader struct {
+func fakeExecServer(t *testing.T, i int, stdinData, stdoutData, stderrData, errorData string, tty bool) http.HandlerFunc {
-	conn *fakeUpgradeConnection
+	// error + stdin + stdout
-	err  error
+	expectedStreams := 3
+	if !tty {
+		// stderr
+		expectedStreams++
 	}
 
-func (u *fakeUpgrader) upgrade(req *client.Request, config *client.Config) (httpstream.Connection, error) {
+	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
-	return u.conn, u.err
+		streamCh := make(chan httpstream.Stream)
+
+		upgrader := spdy.NewResponseUpgrader()
+		conn := upgrader.UpgradeResponse(w, req, func(stream httpstream.Stream) error {
+			streamCh <- stream
+			return nil
+		})
+		// from this point on, we can no longer call methods on w
+		if conn == nil {
+			// The upgrader is responsible for notifying the client of any errors that
+			// occurred during upgrading. All we can do is return here at this point
+			// if we weren't successful in upgrading.
+			return
 		}
+		defer conn.Close()
 
-type fakeUpgradeConnection struct {
+		var errorStream, stdinStream, stdoutStream, stderrStream httpstream.Stream
-	closeCalled bool
+		receivedStreams := 0
-	lock        sync.Mutex
+	WaitForStreams:
-
+		for {
-	stdin                   *fakeUpgradeStream
+			select {
-	stdout                  *fakeUpgradeStream
+			case stream := <-streamCh:
-	stdoutData              string
+				streamType := stream.Headers().Get(api.StreamType)
-	stderr                  *fakeUpgradeStream
+				switch streamType {
-	stderrData              string
-	errorStream             *fakeUpgradeStream
-	errorData               string
-	unexpectedStreamCreated bool
-}
-
-func newFakeUpgradeConnection() *fakeUpgradeConnection {
-	return &fakeUpgradeConnection{}
-}
-
-func (c *fakeUpgradeConnection) CreateStream(headers http.Header) (httpstream.Stream, error) {
-	c.lock.Lock()
-	defer c.lock.Unlock()
-
-	stream := &fakeUpgradeStream{}
-	switch headers.Get(api.StreamType) {
-	case api.StreamTypeStdin:
-		c.stdin = stream
-	case api.StreamTypeStdout:
-		c.stdout = stream
-		stream.data = c.stdoutData
-	case api.StreamTypeStderr:
-		c.stderr = stream
-		stream.data = c.stderrData
 				case api.StreamTypeError:
-		c.errorStream = stream
+					errorStream = stream
-		stream.data = c.errorData
+					receivedStreams++
+				case api.StreamTypeStdin:
+					stdinStream = stream
+					stdinStream.Close()
+					receivedStreams++
+				case api.StreamTypeStdout:
+					stdoutStream = stream
+					receivedStreams++
+				case api.StreamTypeStderr:
+					stderrStream = stream
+					receivedStreams++
 				default:
-		c.unexpectedStreamCreated = true
+					t.Errorf("%d: unexpected stream type: %q", i, streamType)
 				}
 
-	return stream, nil
+				defer stream.Reset()
+
+				if receivedStreams == expectedStreams {
+					break WaitForStreams
+				}
+			}
 		}
 
-func (c *fakeUpgradeConnection) Close() error {
+		if len(errorData) > 0 {
-	c.lock.Lock()
+			fmt.Fprint(errorStream, errorData)
-	defer c.lock.Unlock()
+			errorStream.Close()
-
-	c.closeCalled = true
-	return nil
 		}
 
-func (c *fakeUpgradeConnection) CloseChan() <-chan bool {
+		if len(stdoutData) > 0 {
-	return make(chan bool)
+			fmt.Fprint(stdoutStream, stdoutData)
+			stdoutStream.Close()
 		}
-
+		if len(stderrData) > 0 {
-func (c *fakeUpgradeConnection) SetIdleTimeout(timeout time.Duration) {
+			fmt.Fprint(stderrStream, stderrData)
+			stderrStream.Close()
 		}
-
+		if len(stdinData) > 0 {
-type fakeUpgradeStream struct {
+			data, err := ioutil.ReadAll(stdinStream)
-	readCalled  bool
+			if err != nil {
-	writeCalled bool
+				t.Errorf("%d: error reading stdin stream: %v", i, err)
-	dataWritten []byte
-	closeCalled bool
-	resetCalled bool
-	data        string
-	lock        sync.Mutex
 			}
-
+			if e, a := stdinData, string(data); e != a {
-func (s *fakeUpgradeStream) Read(p []byte) (int, error) {
+				t.Errorf("%d: stdin: expected %q, got %q", i, e, a)
-	s.lock.Lock()
-	defer s.lock.Unlock()
-	s.readCalled = true
-	b := []byte(s.data)
-	n := copy(p, b)
-	return n, io.EOF
 			}
-
-func (s *fakeUpgradeStream) Write(p []byte) (int, error) {
-	s.lock.Lock()
-	defer s.lock.Unlock()
-	s.writeCalled = true
-	s.dataWritten = make([]byte, len(p))
-	copy(s.dataWritten, p)
-	return len(p), io.EOF
 		}
-
+	})
-func (s *fakeUpgradeStream) Close() error {
-	s.lock.Lock()
-	defer s.lock.Unlock()
-	s.closeCalled = true
-	return nil
-}
-
-func (s *fakeUpgradeStream) Reset() error {
-	s.lock.Lock()
-	defer s.lock.Unlock()
-	s.resetCalled = true
-	return nil
-}
-
-func (s *fakeUpgradeStream) Headers() http.Header {
-	s.lock.Lock()
-	defer s.lock.Unlock()
-	return http.Header{}
 }
 
 func TestRequestExecuteRemoteCommand(t *testing.T) {
 	testCases := []struct {
-		Upgrader    *fakeUpgrader
 		Stdin  string
 		Stdout string
 		Stderr string
 		Error  string
 		Tty    bool
-		ShouldError bool
 	}{
 		{
-			Upgrader:    &fakeUpgrader{err: errors.New("bail")},
-			ShouldError: true,
-		},
-		{
-			Upgrader:    &fakeUpgrader{conn: newFakeUpgradeConnection()},
-			Stdin:       "a",
-			Stdout:      "b",
-			Stderr:      "c",
 			Error: "bail",
-			ShouldError: true,
 		},
 		{
-			Upgrader: &fakeUpgrader{conn: newFakeUpgradeConnection()},
 			Stdin:  "a",
 			Stdout: "b",
 			Stderr: "c",
 		},
 		{
-			Upgrader: &fakeUpgrader{conn: newFakeUpgradeConnection()},
 			Stdin:  "a",
 			Stdout: "b",
-			Stderr:   "c",
 			Tty:    true,
 		},
 	}
 
 	for i, testCase := range testCases {
-		if testCase.Error != "" {
+		localOut := &bytes.Buffer{}
-			testCase.Upgrader.conn.errorData = testCase.Error
+		localErr := &bytes.Buffer{}
+
+		server := httptest.NewServer(fakeExecServer(t, i, testCase.Stdin, testCase.Stdout, testCase.Stderr, testCase.Error, testCase.Tty))
+
+		url, _ := url.ParseRequestURI(server.URL)
+		c := client.NewRESTClient(url, "x", nil, -1, -1)
+		req := c.Post().Resource("testing")
+
+		conf := &client.Config{
+			Host: server.URL,
 		}
-		if testCase.Stdout != "" {
+		e := New(req, conf, []string{"ls", "/"}, strings.NewReader(testCase.Stdin), localOut, localErr, testCase.Tty)
-			testCase.Upgrader.conn.stdoutData = testCase.Stdout
+		//e.upgrader = testCase.Upgrader
-		}
-		if testCase.Stderr != "" {
-			testCase.Upgrader.conn.stderrData = testCase.Stderr
-		}
-		var localOut, localErr *bytes.Buffer
-		if testCase.Stdout != "" {
-			localOut = &bytes.Buffer{}
-		}
-		if testCase.Stderr != "" {
-			localErr = &bytes.Buffer{}
-		}
-		e := New(&client.Request{}, &client.Config{}, []string{"ls", "/"}, strings.NewReader(testCase.Stdin), localOut, localErr, testCase.Tty)
-		e.upgrader = testCase.Upgrader
 		err := e.Execute()
 		hasErr := err != nil
-		if hasErr != testCase.ShouldError {
-			t.Fatalf("%d: expected %t, got %t: %v", i, testCase.ShouldError, hasErr, err)
-		}
 
-		conn := testCase.Upgrader.conn
+		if len(testCase.Error) > 0 {
-		if testCase.Error != "" {
+			if !hasErr {
-			if conn.errorStream == nil {
+				t.Errorf("%d: expected an error", i)
-				t.Fatalf("%d: expected error stream creation", i)
+			} else {
-			}
-			if !conn.errorStream.readCalled {
-				t.Fatalf("%d: expected error stream read", i)
-			}
 				if e, a := testCase.Error, err.Error(); !strings.Contains(a, e) {
-				t.Fatalf("%d: expected error stream read '%v', got '%v'", i, e, a)
+					t.Errorf("%d: expected error stream read '%v', got '%v'", i, e, a)
-			}
-			if !conn.errorStream.resetCalled {
-				t.Fatalf("%d: expected error reset", i)
 				}
 			}
 
-		if testCase.ShouldError {
+			server.Close()
 			continue
 		}
 
-		if testCase.Stdin != "" {
+		if hasErr {
-			if conn.stdin == nil {
+			t.Errorf("%d: unexpected error: %v", i, err)
-				t.Fatalf("%d: expected stdin stream creation", i)
+			server.Close()
-			}
+			continue
-			if !conn.stdin.writeCalled {
-				t.Fatalf("%d: expected stdin stream write", i)
-			}
-			if e, a := testCase.Stdin, string(conn.stdin.dataWritten); e != a {
-				t.Fatalf("%d: expected stdin write %v, got %v", i, e, a)
-			}
-			if !conn.stdin.resetCalled {
-				t.Fatalf("%d: expected stdin reset", i)
-			}
 		}
 
-		if testCase.Stdout != "" {
+		if len(testCase.Stdout) > 0 {
-			if conn.stdout == nil {
-				t.Fatalf("%d: expected stdout stream creation", i)
-			}
-			if !conn.stdout.readCalled {
-				t.Fatalf("%d: expected stdout stream read", i)
-			}
 			if e, a := testCase.Stdout, localOut; e != a.String() {
-				t.Fatalf("%d: expected stdout data '%s', got '%s'", i, e, a)
+				t.Errorf("%d: expected stdout data '%s', got '%s'", i, e, a)
-			}
-			if !conn.stdout.resetCalled {
-				t.Fatalf("%d: expected stdout reset", i)
 			}
 		}
 
 		if testCase.Stderr != "" {
-			if testCase.Tty {
-				if conn.stderr != nil {
-					t.Fatalf("%d: unexpected stderr stream creation", i)
-				}
-				if localErr.String() != "" {
-					t.Fatalf("%d: unexpected stderr data '%s'", i, localErr)
-				}
-			} else {
-				if conn.stderr == nil {
-					t.Fatalf("%d: expected stderr stream creation", i)
-				}
-				if !conn.stderr.readCalled {
-					t.Fatalf("%d: expected stderr stream read", i)
-				}
 			if e, a := testCase.Stderr, localErr; e != a.String() {
-					t.Fatalf("%d: expected stderr data '%s', got '%s'", i, e, a)
+				t.Errorf("%d: expected stderr data '%s', got '%s'", i, e, a)
-				}
-				if !conn.stderr.resetCalled {
-					t.Fatalf("%d: expected stderr reset", i)
-				}
 			}
 		}
 
-		if !conn.closeCalled {
+		server.Close()
-			t.Fatalf("%d: expected upgraded connection to get closed", i)
-		}
 	}
 }