fix flaky test TestHandlerStopsWhenContextIsCancelled

aws · Jul 24, 2018 · 56bf075 · 56bf075
1 parent fdb8ca0
commit 56bf075
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/agent/acs/handler/acs_handler_test.go b/agent/acs/handler/acs_handler_test.go
@@ -48,6 +48,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/wsclient/mock"
 	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/gorilla/websocket"
+	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/golang/mock/gomock"
@@ -667,7 +668,7 @@ func TestHandlerStopsWhenContextIsCancelled(t *testing.T) {
 		mockWsClient.EXPECT().Serve().Return(io.EOF),
 		mockWsClient.EXPECT().Serve().Do(func() {
 			cancel()
-		}).Return(io.EOF),
+		}).Return(errors.New("InactiveInstanceException")),
 	)
 	acsSession := session{
 		containerInstanceARN: "myArn",