mutex;

Author: ItamarYuran

cmd/lakectl/cmd/token_cacing_test.go

@@ -195,7 +195,16 @@ func TestTokenExpiredWrite(t *testing.T) {
 		require.Nil(t, token)
 	})
 }
+
+var testMutex sync.Mutex
+
 func resetGlobalState() {
+	testMutex.Lock()
+	defer testMutex.Unlock()
+
+	// Wait longer to ensure all callback goroutines complete
+	time.Sleep(500 * time.Millisecond)
+
 	tokenLoadOnce = sync.Once{}
 	tokenCacheOnce = sync.Once{}
 	tokenSaveOnce = sync.Once{}
@@ -208,7 +217,6 @@ func resetGlobalState() {
 		os.RemoveAll(cacheDir)
 	}
 }
-
 func setupTestHomeDir(t *testing.T) (cleanup func()) {
 	tempDir := t.TempDir()
 	originalHome := os.Getenv("HOME")
@@ -310,12 +318,8 @@ func createSecurityProvider(mockClient *mockExternalLoginClient, initialToken *a
 	)
 }
 
-var testMutex sync.Mutex
-
 func TestLoginOnlyOnce(t *testing.T) {
 	t.Run("login called only once across multiple requests", func(t *testing.T) {
-		testMutex.Lock()
-		defer testMutex.Unlock()
 		resetGlobalState()
 		cleanup := setupTestHomeDir(t)
 		defer cleanup()
@@ -329,26 +333,18 @@ func TestLoginOnlyOnce(t *testing.T) {
 		provider := createSecurityProvider(mockClient, nil, &callbackCount)
 		require.Nil(t, provider.AuthenticationToken)
 
-		// Add debug logging
-		t.Logf("Before Intercept - Token: %v", provider.AuthenticationToken)
-		t.Logf("Before Intercept - Login count: %d", mockClient.getLoginCount())
-
 		req1 := httptest.NewRequest("GET", "http://example.com/api/v1/repositories", nil)
 		err := provider.Intercept(context.Background(), req1)
 
-		t.Logf("After Intercept - Token: %v", provider.AuthenticationToken)
-		t.Logf("After Intercept - Login count: %d", mockClient.getLoginCount())
-		t.Logf("Authorization header: %s", req1.Header.Get("Authorization"))
-
 		require.NoError(t, err)
 		require.Equal(t, "Bearer cached-token", req1.Header.Get("Authorization"))
 		require.Equal(t, int64(1), mockClient.getLoginCount())
+
+		// Wait for callback goroutine to complete
+		time.Sleep(200 * time.Millisecond)
 	})
-	// }
-	// func TestNoLoginWhenTokenIsGiven(t *testing.T) {
+
 	t.Run("no login performed when valid token provided initially", func(t *testing.T) {
-		testMutex.Lock()
-		defer testMutex.Unlock()
 		resetGlobalState()
 		cleanup := setupTestHomeDir(t)
 		defer cleanup()
@@ -367,20 +363,17 @@ func TestLoginOnlyOnce(t *testing.T) {
 		var callbackCount int64
 		provider := createSecurityProvider(mockClient, existingToken, &callbackCount)
 
-		// Request should use existing token, no login
 		req := httptest.NewRequest("GET", "http://example.com/api/v1/repositories", nil)
 		err := provider.Intercept(context.Background(), req)
 		require.NoError(t, err)
 		require.Equal(t, "Bearer pre-existing-token", req.Header.Get("Authorization"))
-		require.Equal(t, int64(0), mockClient.getLoginCount()) // No login called
-		time.Sleep(time.Second)
-		require.Equal(t, int64(0), atomic.LoadInt64(&callbackCount)) // No callback called
+		require.Equal(t, int64(0), mockClient.getLoginCount())
+
+		time.Sleep(200 * time.Millisecond)
+		require.Equal(t, int64(0), atomic.LoadInt64(&callbackCount))
 	})
-	// }
-	// func TestRealInterceptWithGlobalCache2(t *testing.T) {
+
 	t.Run("handles login failure gracefully", func(t *testing.T) {
-		testMutex.Lock()
-		defer testMutex.Unlock()
 		resetGlobalState()
 		cleanup := setupTestHomeDir(t)
 		defer cleanup()
@@ -393,20 +386,17 @@ func TestLoginOnlyOnce(t *testing.T) {
 		var callbackCount int64
 		provider := createSecurityProvider(mockClient, nil, &callbackCount)
 
-		// Request should fail due to login failure
 		req := httptest.NewRequest("GET", "http://example.com/api/v1/repositories", nil)
 		err := provider.Intercept(context.Background(), req)
 		require.ErrorIs(t, err, errMockLoginFailed)
 		require.Empty(t, req.Header.Get("Authorization"))
-		require.Equal(t, int64(1), mockClient.getLoginCount()) // Login attempted once
-		time.Sleep(time.Second)
-		require.Equal(t, int64(0), atomic.LoadInt64(&callbackCount)) // No callback on failure
+		require.Equal(t, int64(1), mockClient.getLoginCount())
+
+		time.Sleep(200 * time.Millisecond)
+		require.Equal(t, int64(0), atomic.LoadInt64(&callbackCount))
 	})
-	// }
-	// func TestRealInterceptWithGlobalCache5(t *testing.T) {
+
 	t.Run("token cached via callback and reused after provider recreation", func(t *testing.T) {
-		testMutex.Lock()
-		defer testMutex.Unlock()
 		resetGlobalState()
 		cleanup := setupTestHomeDir(t)
 		defer cleanup()
@@ -419,34 +409,31 @@ func TestLoginOnlyOnce(t *testing.T) {
 		var callbackCount int64
 		provider1 := createSecurityProvider(mockClient, nil, &callbackCount)
 
-		// First request - should trigger login and cache via callback
 		req1 := httptest.NewRequest("GET", "http://example.com/api/v1/repositories", nil)
 		err := provider1.Intercept(context.Background(), req1)
 		require.NoError(t, err)
 		require.Equal(t, "Bearer callback-cached-token", req1.Header.Get("Authorization"))
 		require.Equal(t, int64(1), mockClient.getLoginCount())
-		time.Sleep(time.Second)
+
+		time.Sleep(200 * time.Millisecond)
 		require.Equal(t, int64(1), atomic.LoadInt64(&callbackCount))
 
-		// Verify token was saved to global cache
 		require.NotNil(t, cachedToken)
 		require.Equal(t, "callback-cached-token", cachedToken.Token)
 
-		// Create new provider instance (simulating new command execution)
 		mockClient.resetLoginCount()
 		callbackCount = 0
 
-		// Get token from cache using global function
 		cachedTokenFromFile := getTokenOnce()
 		provider2 := createSecurityProvider(mockClient, cachedTokenFromFile, &callbackCount)
 
-		// Second request with new provider should use cached token from file
 		req2 := httptest.NewRequest("GET", "http://example.com/api/v1/repositories", nil)
 		err = provider2.Intercept(context.Background(), req2)
 		require.NoError(t, err)
 		require.Equal(t, "Bearer callback-cached-token", req2.Header.Get("Authorization"))
-		require.Equal(t, int64(0), mockClient.getLoginCount()) // No new login, used cache
-		time.Sleep(time.Second)
-		require.Equal(t, int64(0), atomic.LoadInt64(&callbackCount)) // No new callback
+		require.Equal(t, int64(0), mockClient.getLoginCount())
+
+		time.Sleep(200 * time.Millisecond)
+		require.Equal(t, int64(0), atomic.LoadInt64(&callbackCount))
 	})
 }