fix: flaky test in TestProxySocketsAdminFacade (#1474)

* fix: flaky test in TestProxySocketsAdminFacade

The test was incorrectly cancelling the context. When the proxySockets function encounters an error like a failure to connect to a Juju controller it returns with an error. Otherwise it returns a nil error when the context is cancelled. Here we were expecting the former for one of the tests cases but sometimes encountering the latter in a race condition.

* chore: remove expectProxyError bool from test

internal/rpc/apiproxy_test.go

@@ -70,7 +70,7 @@ func TestProxySocketsAdminFacade(t *testing.T) {
 		expectedClientResponse    *message
 		expectedControllerMessage *message
 		oauthAuthenticatorError   error
-		expectConnectTofail       bool
+		expectedProxyError        string
 	}{{
 		about: "login device call - client gets response with both user code and verification uri",
 		messageToSend: message{
@@ -232,18 +232,21 @@ func TestProxySocketsAdminFacade(t *testing.T) {
 		},
 		oauthAuthenticatorError: errors.E(errors.CodeUnauthorized),
 	}, {
-		about:               "connection to controller fails",
-		expectConnectTofail: true,
+		about: "connection to controller fails",
 		expectedClientResponse: &message{
 			Error: "controller connection error",
 		},
+		expectedProxyError: "failed to connect to controller: controller connection error",
 	}}
 
 	for _, test := range tests {
 		t.Run(test.about, func(t *testing.T) {
+			proxyError := test.expectedProxyError != ""
+
 			ctx := context.Background()
 			ctx, cancelFunc := context.WithCancel(ctx)
 			defer cancelFunc()
+
 			clientWebsocket := newMockWebsocketConnection(10)
 			controllerWebsocket := newMockWebsocketConnection(10)
 			loginSvc := &mockLoginService{
@@ -257,7 +260,7 @@ func TestProxySocketsAdminFacade(t *testing.T) {
 				ConnClient: clientWebsocket,
 				TokenGen:   &mockTokenGenerator{},
 				ConnectController: func(ctx context.Context) (rpc.WebsocketConnectionWithMetadata, error) {
-					if test.expectConnectTofail {
+					if proxyError {
 						return rpc.WebsocketConnectionWithMetadata{}, goerr.New("controller connection error")
 					}
 					return rpc.WebsocketConnectionWithMetadata{
@@ -275,8 +278,8 @@ func TestProxySocketsAdminFacade(t *testing.T) {
 			go func() {
 				defer wg.Done()
 				err = rpc.ProxySockets(ctx, helpers)
-				if test.expectConnectTofail {
-					c.Assert(err, qt.ErrorMatches, "failed to connect to controller: controller connection error")
+				if proxyError {
+					c.Assert(err, qt.ErrorMatches, test.expectedProxyError)
 				} else {
 					c.Assert(err, qt.ErrorMatches, "Context cancelled")
 				}
@@ -300,7 +303,9 @@ func TestProxySocketsAdminFacade(t *testing.T) {
 					c.Fatal("timed out waiting for response")
 				}
 			}
-			cancelFunc()
+			if !proxyError {
+				cancelFunc()
+			}
 			wg.Wait()
 			t.Logf("completed test %s", t.Name())
 		})