add auth_token_async_test.go

Author: quartzmo

auth/auth_test.go

@@ -22,7 +22,6 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strings"
-	"sync"
 	"testing"
 	"time"
 
@@ -106,9 +105,9 @@ func TestError_Temporary(t *testing.T) {
 
 func TestToken_MetadataString(t *testing.T) {
 	cases := []struct {
-		name string
+		name     string
 		metadata map[string]interface{}
-		want string
+		want     string
 	}{
 		{
 			name: "nil metadata",
@@ -143,10 +142,10 @@ func TestToken_isValidWithEarlyExpiry(t *testing.T) {
 	defer func() { timeNow = time.Now }()
 
 	cases := []struct {
-		name string
-		tok *Token
+		name   string
+		tok    *Token
 		expiry time.Duration
-		want bool
+		want   bool
 	}{
 		{name: "4 minutes", tok: &Token{Expiry: now.Add(4 * 60 * time.Second)}, expiry: defaultExpiryDelta, want: true},
 		{name: "3 minutes and 45 seconds", tok: &Token{Expiry: now.Add(defaultExpiryDelta)}, expiry: defaultExpiryDelta, want: true},
@@ -170,12 +169,12 @@ func TestError_Error(t *testing.T) {
 	tests := []struct {
 		name string
 
-		Response *http.Response
-		Body []byte
-		Err error
-		code string
+		Response    *http.Response
+		Body        []byte
+		Err         error
+		code        string
 		description string
-		uri string
+		uri         string
 
 		want string
 	}{
@@ -188,22 +187,22 @@ func TestError_Error(t *testing.T) {
 			want: "auth: cannot fetch token: 418\nResponse: I'm a teapot",
 		},
 		{
-			name: "from query",
-			code: fmt.Sprint(http.StatusTeapot),
+			name:        "from query",
+			code:        fmt.Sprint(http.StatusTeapot),
 			description: "I'm a teapot",
-			uri: "somewhere",
-			want: "auth: \"418\" \"I'm a teapot\" \"somewhere\"",
+			uri:         "somewhere",
+			want:        "auth: \"418\" \"I'm a teapot\" \"somewhere\"",
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			r := &Error{
-				Response: tt.Response,
-				Body: tt.Body,
-				Err: tt.Err,
-				code: tt.code,
+				Response:    tt.Response,
+				Body:        tt.Body,
+				Err:         tt.Err,
+				code:        tt.code,
 				description: tt.description,
-				uri: tt.uri,
+				uri:         tt.uri,
 			}
 			if got := r.Error(); got != tt.want {
 				t.Errorf("Error.Error() = %v, want %v", got, tt.want)
@@ -225,9 +224,9 @@ func TestNew2LOTokenProvider_JSONResponse(t *testing.T) {
 	defer ts.Close()
 
 	opts := &Options2LO{
-		Email: "aaa@example.com",
+		Email:      "aaa@example.com",
 		PrivateKey: fakePrivateKey,
-		TokenURL: ts.URL,
+		TokenURL:   ts.URL,
 	}
 	tp, err := New2LOTokenProvider(opts)
 	if err != nil {
@@ -263,9 +262,9 @@ func TestNew2LOTokenProvider_BadResponse(t *testing.T) {
 	defer ts.Close()
 
 	opts := &Options2LO{
-		Email: "aaa@example.com",
+		Email:      "aaa@example.com",
 		PrivateKey: fakePrivateKey,
-		TokenURL: ts.URL,
+		TokenURL:   ts.URL,
 	}
 	tp, err := New2LOTokenProvider(opts)
 	if err != nil {
@@ -300,9 +299,9 @@ func TestNew2LOTokenProvider_BadResponseType(t *testing.T) {
 	}))
 	defer ts.Close()
 	opts := &Options2LO{
-		Email: "aaa@example.com",
+		Email:      "aaa@example.com",
 		PrivateKey: fakePrivateKey,
-		TokenURL: ts.URL,
+		TokenURL:   ts.URL,
 	}
 	tp, err := New2LOTokenProvider(opts)
 	if err != nil {
@@ -334,10 +333,10 @@ func TestNew2LOTokenProvider_Assertion(t *testing.T) {
 	defer ts.Close()
 
 	opts := &Options2LO{
-		Email: "aaa@example.com",
-		PrivateKey: fakePrivateKey,
+		Email:        "aaa@example.com",
+		PrivateKey:   fakePrivateKey,
 		PrivateKeyID: "ABCDEFGHIJKLMNOPQRSTUVWXYZ",
-		TokenURL: ts.URL,
+		TokenURL:     ts.URL,
 	}
 
 	tp, err := New2LOTokenProvider(opts)
@@ -365,8 +364,8 @@ func TestNew2LOTokenProvider_Assertion(t *testing.T) {
 
 	want := jwt.Header{
 		Algorithm: "RS256",
-		Type: "JWT",
-		KeyID: "ABCDEFGHIJKLMNOPQRSTUVWXYZ",
+		Type:      "JWT",
+		KeyID:     "ABCDEFGHIJKLMNOPQRSTUVWXYZ",
 	}
 	if got != want {
 		t.Errorf("access token header = %q; want %q", got, want)
@@ -391,23 +390,23 @@ func TestNew2LOTokenProvider_AssertionPayload(t *testing.T) {
 
 	for _, opts := range []*Options2LO{
 		{
-			Email: "aaa1@example.com",
-			PrivateKey: fakePrivateKey,
+			Email:        "aaa1@example.com",
+			PrivateKey:   fakePrivateKey,
 			PrivateKeyID: "ABCDEFGHIJKLMNOPQRSTUVWXYZ",
-			TokenURL: ts.URL,
+			TokenURL:     ts.URL,
 		},
 		{
-			Email: "aaa2@example.com",
-			PrivateKey: fakePrivateKey,
+			Email:        "aaa2@example.com",
+			PrivateKey:   fakePrivateKey,
 			PrivateKeyID: "ABCDEFGHIJKLMNOPQRSTUVWXYZ",
-			TokenURL: ts.URL,
-			Audience: "https://example.com",
+			TokenURL:     ts.URL,
+			Audience:     "https://example.com",
 		},
 		{
-			Email: "aaa2@example.com",
-			PrivateKey: fakePrivateKey,
+			Email:        "aaa2@example.com",
+			PrivateKey:   fakePrivateKey,
 			PrivateKeyID: "ABCDEFGHIJKLMNOPQRSTUVWXYZ",
-			TokenURL: ts.URL,
+			TokenURL:     ts.URL,
 			PrivateClaims: map[string]interface{}{
 				"private0": "claim0",
 				"private1": "claim1",
@@ -479,9 +478,9 @@ func TestNew2LOTokenProvider_TokenError(t *testing.T) {
 	defer ts.Close()
 
 	opts := &Options2LO{
-		Email: "aaa@example.com",
+		Email:      "aaa@example.com",
 		PrivateKey: fakePrivateKey,
-		TokenURL: ts.URL,
+		TokenURL:   ts.URL,
 	}
 
 	tp, err := New2LOTokenProvider(opts)
@@ -514,20 +513,20 @@ func TestNew2LOTokenProvider_Validate(t *testing.T) {
 			name: "missing email",
 			opts: &Options2LO{
 				PrivateKey: []byte("key"),
-				TokenURL: "url",
+				TokenURL:   "url",
 			},
 		},
 		{
 			name: "missing key",
 			opts: &Options2LO{
-				Email: "email",
+				Email:    "email",
 				TokenURL: "url",
 			},
 		},
 		{
 			name: "missing URL",
 			opts: &Options2LO{
-				Email: "email",
+				Email:      "email",
 				PrivateKey: []byte("key"),
 			},
 		},
@@ -616,25 +615,25 @@ func TestComputeTokenProvider_NonBlockingRefresh(t *testing.T) {
 
 func TestComputeTokenProvider_BlockingRefresh(t *testing.T) {
 	tests := []struct {
-		name string
+		name               string
 		disableAutoRefresh bool
-		want1 string
-		want2 string
-		wantState2 tokenState
+		want1              string
+		want2              string
+		wantState2         tokenState
 	}{
 		{
-			name: "disableAutoRefresh",
+			name:               "disableAutoRefresh",
 			disableAutoRefresh: true,
-			want1: "1",
-			want2: "1",
+			want1:              "1",
+			want2:              "1",
 			// Because token "count" does not increase, it will always be stale.
 			wantState2: stale,
 		},
 		{
-			name: "autoRefresh",
+			name:               "autoRefresh",
 			disableAutoRefresh: false,
-			want1: "1",
-			want2: "2",
+			want1:              "1",
+			want2:              "2",
 			// As token "count" increases to 2, it transitions to fresh.
 			wantState2: fresh,
 		},
@@ -647,7 +646,7 @@ func TestComputeTokenProvider_BlockingRefresh(t *testing.T) {
 			defer func() { timeNow = time.Now }()
 			tp := NewCachedTokenProvider(&countingTestProvider{count: 1}, &CachedTokenProviderOptions{
 				DisableAsyncRefresh: true,
-				DisableAutoRefresh: tt.disableAutoRefresh,
+				DisableAutoRefresh:  tt.disableAutoRefresh,
 				// EarlyTokenRefresh ensures that token with early expiry just less than 2 seconds before now is already stale.
 				ExpireEarly: 1990 * time.Millisecond,
 			})
@@ -681,97 +680,3 @@ func TestComputeTokenProvider_BlockingRefresh(t *testing.T) {
 		})
 	}
 }
-
-type controllableTokenProvider struct {
-	mu sync.Mutex
-	count int
-	tok *Token
-	err error
-	block chan struct{}
-}
-
-func (p *controllableTokenProvider) Token(ctx context.Context) (*Token, error) {
-	if ch := p.getBlockChan(); ch != nil {
-		<-ch
-	}
-	p.mu.Lock()
-	defer p.mu.Unlock()
-	p.count++
-	return p.tok, p.err
-}
-
-func (p *controllableTokenProvider) getBlockChan() chan struct{} {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-	return p.block
-}
-
-func (p *controllableTokenProvider) setBlockChan(ch chan struct{}) {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-	p.block = ch
-}
-
-func (p *controllableTokenProvider) getCount() int {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-	return p.count
-}
-
-func TestCachedTokenProvider_TokenAsyncRace(t *testing.T) {
-	now := time.Now()
-	timeNow = func() time.Time { return now }
-	defer func() { timeNow = time.Now }()
-
-	tp := &controllableTokenProvider{}
-	ctp := NewCachedTokenProvider(tp, &CachedTokenProviderOptions{
-		ExpireEarly: 2 * time.Second,
-	}).(*cachedTokenProvider)
-
-	// 1. Cache a stale token.
-	tp.tok = &Token{Value: "initial", Expiry: now.Add(1 * time.Second)}
-	if _, err := ctp.Token(context.Background()); err != nil {
-		t.Fatalf("initial Token() failed: %v", err)
-	}
-	if got, want := tp.getCount(), 1; got != want {
-		t.Fatalf("tp.count = %d; want %d", got, want)
-	}
-	if got, want := ctp.tokenState(), stale; got != want {
-		t.Fatalf("tokenState = %v; want %v", got, want)
-	}
-
-	// 2. Setup for refresh.
-	tp.setBlockChan(make(chan struct{}))
-	tp.tok = &Token{Value: "refreshed", Expiry: now.Add(1 * time.Hour)}
-
-	// 3. Concurrently call Token to trigger async refresh.
-	var wg sync.WaitGroup
-	for i := 0; i < 10; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			ctp.Token(context.Background())
-		}()
-	}
-
-	// 4. Unblock refresh and wait for all goroutines to finish.
-	time.Sleep(100 * time.Millisecond) // give time for goroutines to run
-	close(tp.getBlockChan())
-	wg.Wait()
-	time.Sleep(100 * time.Millisecond) // give time for async refresh to complete
-
-	// 5. Check results.
-	if got, want := tp.getCount(), 2; got != want {
-		t.Errorf("tp.count = %d; want %d", got, want)
-	}
-	if got, want := ctp.tokenState(), fresh; got != want {
-		t.Errorf("tokenState = %v; want %v", got, want)
-	}
-	tok, err := ctp.Token(context.Background())
-	if err != nil {
-		t.Fatal(err)
-	}
-	if got, want := tok.Value, "refreshed"; got != want {
-		t.Errorf("tok.Value = %q; want %q", got, want)
-	}
-}