Review: improve Azure handler.

- replace `context.WithTimeout` with `s.Clock.After` to avoid `time.Sleep` in tests.

lib/srv/app/azure/handler.go

@@ -287,19 +287,23 @@ const getTokenTimeout = time.Second * 5
 func (s *handler) getToken(ctx context.Context, managedIdentity string, scope string) (*azcore.AccessToken, error) {
 	key := cacheKey{managedIdentity, scope}
 
-	timeoutCtx, cancel := context.WithTimeout(ctx, getTokenTimeout)
+	cancelCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	token, err := utils.FnCacheGet(timeoutCtx, s.tokenCache, key, func(ctx context.Context) (*azcore.AccessToken, error) {
-		return s.getAccessToken(ctx, managedIdentity, scope)
-	})
-
-	if err != nil {
-		if timeoutCtx.Err() == err {
-			return nil, trace.Wrap(err, "timeout waiting for access token for %v", getTokenTimeout)
-		}
-		return nil, trace.Wrap(err)
+	var tokenResult *azcore.AccessToken
+	var errorResult error
+
+	go func() {
+		tokenResult, errorResult = utils.FnCacheGet(cancelCtx, s.tokenCache, key, func(ctx context.Context) (*azcore.AccessToken, error) {
+			return s.getAccessToken(ctx, managedIdentity, scope)
+		})
+		cancel()
+	}()
+
+	select {
+	case <-s.Clock.After(getTokenTimeout):
+		return nil, trace.Wrap(context.DeadlineExceeded, "timeout waiting for access token for %v", getTokenTimeout)
+	case <-cancelCtx.Done():
+		return tokenResult, errorResult
 	}
-
-	return token, nil
 }

lib/srv/app/azure/handler_test.go

@@ -28,7 +28,7 @@ import (
 func TestForwarder_getToken(t *testing.T) {
 	t.Parallel()
 
-	tests := []struct {
+	type testCase struct {
 		name string
 
 		config HandlerConfig
@@ -38,7 +38,11 @@ func TestForwarder_getToken(t *testing.T) {
 
 		wantToken *azcore.AccessToken
 		checkErr  require.ErrorAssertionFunc
-	}{
+	}
+
+	var tests []testCase
+
+	tests = []testCase{
 		{
 			name: "base case",
 			config: HandlerConfig{
@@ -60,8 +64,18 @@ func TestForwarder_getToken(t *testing.T) {
 		{
 			name: "timeout",
 			config: HandlerConfig{
+				Clock: clockwork.NewFakeClock(),
 				getAccessToken: func(ctx context.Context, managedIdentity string, scope string) (*azcore.AccessToken, error) {
-					time.Sleep(getTokenTimeout * 2)
+					// find the fake clock from above
+					var clock clockwork.FakeClock
+					for _, test := range tests {
+						if test.name == "timeout" {
+							clock = test.config.Clock.(clockwork.FakeClock)
+						}
+					}
+
+					clock.Advance(getTokenTimeout)
+					clock.Sleep(getTokenTimeout * 2)
 					return &azcore.AccessToken{Token: "foobar"}, nil
 				},
 			},