From c847d0acc06d40f2fb1094dfdb3974f70c36c95f Mon Sep 17 00:00:00 2001
From: Tomasz Maczukin <tomasz@gitlab.com>
Date: Tue, 16 Oct 2018 18:14:22 +0000
Subject: [PATCH] Merge branch
 '3639-follow-up-from-resolve-add-docker-support-for-interactive-web-terminal'
 into 'master'

Resolve "Follow-up from Add docker support for interactive web terminal"

Closes #3639

See merge request gitlab-org/gitlab-runner!1049
---
 executors/docker/terminal_test.go | 60 ++++++++++---------------------
 session/server_test.go            |  5 +--
 session/session_test.go           | 43 +++++++++++++---------
 3 files changed, 48 insertions(+), 60 deletions(-)

diff --git a/executors/docker/terminal_test.go b/executors/docker/terminal_test.go
index 0a217d035..b3139ee25 100644
--- a/executors/docker/terminal_test.go
+++ b/executors/docker/terminal_test.go
@@ -109,35 +109,6 @@ func TestInteractiveTerminal(t *testing.T) {
 	assert.Contains(t, string(tty), "Linux")
 }
 
-func TestCommandExecutor_Connect_Timeout(t *testing.T) {
-	c := &docker_helpers.MockClient{}
-
-	s := commandExecutor{
-		executor: executor{
-			AbstractExecutor: executors.AbstractExecutor{
-				Context: context.Background(),
-				BuildShell: &common.ShellConfiguration{
-					DockerCommand: []string{"/bin/sh"},
-				},
-			},
-			client: c,
-		},
-		buildContainer: &types.ContainerJSON{
-			ContainerJSONBase: &types.ContainerJSONBase{
-				ID: "1234",
-			},
-		},
-	}
-	c.On("ContainerInspect", s.Context, "1234").Return(types.ContainerJSON{
-		ContainerJSONBase: &types.ContainerJSONBase{
-			State: &types.ContainerState{
-				Running: false,
-			},
-		},
-	}, nil)
-
-}
-
 func TestCommandExecutor_Connect(t *testing.T) {
 	tests := []struct {
 		name                  string
@@ -166,7 +137,7 @@ func TestCommandExecutor_Connect(t *testing.T) {
 			expectedErr:           errors.New("container not found"),
 		},
 		{
-			name: "Not build container",
+			name: "No build container",
 			buildContainerRunning: false,
 			hasBuildContainer:     false,
 			expectedErr:           buildContainerTerminalTimeout{},
@@ -176,6 +147,7 @@ func TestCommandExecutor_Connect(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			c := &docker_helpers.MockClient{}
+			defer c.AssertExpectations(t)
 
 			s := commandExecutor{
 				executor: executor{
@@ -195,15 +167,15 @@ func TestCommandExecutor_Connect(t *testing.T) {
 						ID: "1234",
 					},
 				}
-			}
 
-			c.On("ContainerInspect", s.Context, "1234").Return(types.ContainerJSON{
-				ContainerJSONBase: &types.ContainerJSONBase{
-					State: &types.ContainerState{
-						Running: test.buildContainerRunning,
+				c.On("ContainerInspect", s.Context, "1234").Return(types.ContainerJSON{
+					ContainerJSONBase: &types.ContainerJSONBase{
+						State: &types.ContainerState{
+							Running: test.buildContainerRunning,
+						},
 					},
-				},
-			}, test.containerInspectErr)
+				}, test.containerInspectErr)
+			}
 
 			conn, err := s.Connect()
 
@@ -242,6 +214,7 @@ func TestTerminalConn_FailToStart(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			c := &docker_helpers.MockClient{}
+			defer c.AssertExpectations(t)
 
 			s := commandExecutor{
 				executor: executor{
@@ -271,12 +244,14 @@ func TestTerminalConn_FailToStart(t *testing.T) {
 			c.On("ContainerExecCreate", mock.Anything, mock.Anything, mock.Anything).Return(
 				types.IDResponse{},
 				test.containerExecCreateErr,
-			)
+			).Once()
 
-			c.On("ContainerExecAttach", mock.Anything, mock.Anything, mock.Anything).Return(
-				types.HijackedResponse{},
-				test.containerExecAttachErr,
-			)
+			if test.containerExecCreateErr == nil {
+				c.On("ContainerExecAttach", mock.Anything, mock.Anything, mock.Anything).Return(
+					types.HijackedResponse{},
+					test.containerExecAttachErr,
+				).Once()
+			}
 
 			conn, err := s.Connect()
 			require.NoError(t, err)
@@ -337,6 +312,7 @@ func (nopConn) SetWriteDeadline(t time.Time) error {
 
 func TestTerminalConn_Start(t *testing.T) {
 	c := &docker_helpers.MockClient{}
+	defer c.AssertExpectations(t)
 
 	s := commandExecutor{
 		executor: executor{
diff --git a/session/server_test.go b/session/server_test.go
index cd77e5972..09eabfe7f 100644
--- a/session/server_test.go
+++ b/session/server_test.go
@@ -127,9 +127,10 @@ func TestFailedToGenerateCertificate(t *testing.T) {
 		ListenAddress: "127.0.0.1:0",
 	}
 
-	m := certificate.MockGenerator{}
+	m := new(certificate.MockGenerator)
+	defer m.AssertExpectations(t)
 	m.On("Generate", mock.Anything).Return(tls.Certificate{}, []byte{}, errors.New("something went wrong"))
 
-	_, err := NewServer(cfg, nil, &m, fakeSessionFinder)
+	_, err := NewServer(cfg, nil, m, fakeSessionFinder)
 	assert.Error(t, err, "something went wrong")
 }
diff --git a/session/session_test.go b/session/session_test.go
index c0034fe16..b090ed407 100644
--- a/session/session_test.go
+++ b/session/session_test.go
@@ -64,15 +64,23 @@ func TestExec(t *testing.T) {
 			require.NoError(t, err)
 			session.Token = "validToken"
 
-			mockTerminalConn := terminal.MockConn{}
-			mockTerminalConn.On("Start", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Once()
-			mockTerminalConn.On("Close").Return(nil).Once()
+			mockTerminalConn := new(terminal.MockConn)
+			defer mockTerminalConn.AssertExpectations(t)
 
-			mockTerminal := terminal.MockInteractiveTerminal{}
-			mockTerminal.On("Connect").Return(&mockTerminalConn, c.connectionErr).Once()
+			if c.connectionErr == nil && c.authorization == "validToken" && c.isWebsocketUpgrade {
+				mockTerminalConn.On("Start", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Once()
+				mockTerminalConn.On("Close").Return(nil).Once()
+			}
+
+			mockTerminal := new(terminal.MockInteractiveTerminal)
+			defer mockTerminal.AssertExpectations(t)
+
+			if c.authorization == "validToken" && c.isWebsocketUpgrade {
+				mockTerminal.On("Connect").Return(mockTerminalConn, c.connectionErr).Once()
+			}
 
 			if c.attachTerminal {
-				session.SetInteractiveTerminal(&mockTerminal)
+				session.SetInteractiveTerminal(mockTerminal)
 			}
 
 			req := httptest.NewRequest(http.MethodPost, session.Endpoint+"/exec", nil)
@@ -99,14 +107,14 @@ func TestDoNotAllowMultipleConnections(t *testing.T) {
 	require.NoError(t, err)
 	session.Token = "validToken"
 
-	mockTerminalConn := terminal.MockConn{}
-	mockTerminalConn.On("Start", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Once()
-	mockTerminalConn.On("Close").Return().Once()
+	mockTerminalConn := new(terminal.MockConn)
+	defer mockTerminalConn.AssertExpectations(t)
 
-	mockTerminal := terminal.MockInteractiveTerminal{}
-	mockTerminal.On("Connect").Return(&mockTerminalConn, nil).Once()
+	mockTerminal := new(terminal.MockInteractiveTerminal)
+	defer mockTerminal.AssertExpectations(t)
+	mockTerminal.On("Connect").Return(mockTerminalConn, nil).Once()
 
-	session.SetInteractiveTerminal(&mockTerminal)
+	session.SetInteractiveTerminal(mockTerminal)
 
 	// Simulating another connection has already started.
 	conn, err := session.newTerminalConn()
@@ -141,10 +149,11 @@ func TestKill(t *testing.T) {
 	err = sess.Kill()
 	assert.NoError(t, err)
 
-	mockConn := terminal.MockConn{}
+	mockConn := new(terminal.MockConn)
+	defer mockConn.AssertExpectations(t)
 	mockConn.On("Close").Return(nil).Once()
 
-	sess.terminalConn = &mockConn
+	sess.terminalConn = mockConn
 
 	err = sess.Kill()
 	assert.NoError(t, err)
@@ -155,10 +164,11 @@ func TestKillFailedToClose(t *testing.T) {
 	sess, err := NewSession(nil)
 	require.NoError(t, err)
 
-	mockConn := terminal.MockConn{}
+	mockConn := new(terminal.MockConn)
+	defer mockConn.AssertExpectations(t)
 	mockConn.On("Close").Return(errors.New("some error")).Once()
 
-	sess.terminalConn = &mockConn
+	sess.terminalConn = mockConn
 
 	err = sess.Kill()
 	assert.Error(t, err)
@@ -184,6 +194,7 @@ func TestCloseTerminalConn(t *testing.T) {
 	}
 
 	mockConn := new(terminal.MockConn)
+	defer mockConn.AssertExpectations(t)
 	mockConn.On("Close").Return(nil).Once()
 
 	sess, err := NewSession(nil)
-- 
GitLab