improve concurrent test

Author: ndyakov

internal/mocks/mocks.go

@@ -133,6 +133,7 @@ var _ net.Error = (*Error)(nil)
 type TokenListener struct {
 	// Mock implementation of the TokenManagerListener interface
 	mock.Mock
+	Id int32
 }
 
 func (m *TokenListener) OnTokenNext(token *token.Token) {

manager/token_manager.go

@@ -114,7 +114,7 @@ func NewTokenManager(idp shared.IdentityProvider, options TokenManagerOptions) (
 	return &entraidTokenManager{
 		idp:                            idp,
 		token:                          nil,
-		closedChan:                     make(chan struct{}),
+		closedChan:                     nil,
 		expirationRefreshRatio:         options.ExpirationRefreshRatio,
 		lowerRefreshBoundMs:            options.LowerRefreshBoundMs,
 		lowerBoundDuration:             time.Duration(options.LowerRefreshBoundMs) * time.Millisecond,
@@ -217,13 +217,12 @@ func (e *entraidTokenManager) Start(listener TokenListener) (CancelFunc, error)
 	if e.listener != nil {
 		return nil, ErrTokenManagerAlreadyStarted
 	}
-	e.listener = listener
+
 	if e.closedChan != nil && !internal.IsClosed(e.closedChan) {
 		// there is a hanging goroutine that is waiting for the closedChan to be closed
 		// if the closedChan is not nil and not closed, close it
 		close(e.closedChan)
 	}
-	e.closedChan = make(chan struct{})
 
 	t, err := e.GetToken(true)
 	if err != nil {
@@ -233,6 +232,9 @@ func (e *entraidTokenManager) Start(listener TokenListener) (CancelFunc, error)
 
 	go listener.OnTokenNext(t)
 
+	e.closedChan = make(chan struct{})
+	e.listener = listener
+
 	go func(listener TokenListener, closed <-chan struct{}) {
 		maxDelay := time.Duration(e.retryOptions.MaxDelayMs) * time.Millisecond
 		initialDelay := time.Duration(e.retryOptions.InitialDelayMs) * time.Millisecond
@@ -311,10 +313,7 @@ func (e *entraidTokenManager) Close() error {
 	if e.closedChan == nil || e.listener == nil {
 		return ErrTokenManagerAlreadyCanceled
 	}
-
-	if e.listener != nil {
-		e.listener = nil
-	}
+	e.listener = nil
 	close(e.closedChan)
 
 	return nil

manager/token_manager_test.go

@@ -367,37 +367,47 @@ func TestTokenManager_Start(t *testing.T) {
 		assert.NoError(t, err)
 
 		assert.NotPanics(t, func() {
-			var last int32
+			last := &atomic.Int32{}
 			wg := &sync.WaitGroup{}
 
 			idp.On("RequestToken").Return(rawResponse, nil)
 			mParser.On("ParseResponse", rawResponse).Return(testTokenValid, nil)
 			listener.On("OnTokenNext", testTokenValid).Return()
-			numExecutions := 50000
-			for i := 0; i < numExecutions; i++ {
+			numExecutions := int32(50000)
+			for i := int32(0); i < numExecutions; i++ {
 				wg.Add(1)
-				go func(num int) {
+				go func(num int32) {
 					defer wg.Done()
-					time.Sleep(time.Duration(int64(rand.Intn(100)) * int64(time.Millisecond)))
 					var err error
+					time.Sleep(time.Duration(int64(rand.Intn(1000)+(300-int(num)/2)) * int64(time.Millisecond)))
+					last.Store(num)
 					if num%2 == 0 {
 						err = tokenManager.Close()
 					} else {
-						_, err = tokenManager.Start(listener)
+						l := &mockTokenListener{Id: num}
+						l.On("OnTokenNext", testTokenValid).Return()
+						_, err = tokenManager.Start(l)
 					}
-					if err == nil {
-						atomic.StoreInt32(&last, int32(num))
-					} else if err != ErrTokenManagerAlreadyStarted && err != ErrTokenManagerAlreadyCanceled {
-						log.Printf("Error: %v", err)
-						t.Fail()
+					if err != nil {
+						if err != ErrTokenManagerAlreadyCanceled && err != ErrTokenManagerAlreadyStarted {
+							// this is un unexpected error, fail the test
+							assert.Error(t, err)
+						}
 					}
 				}(i)
 			}
 			wg.Wait()
-			lastExecution := atomic.LoadInt32(&last)
+			lastExecution := last.Load()
 			if lastExecution%2 == 0 {
+				if tm.listener != nil {
+					l := tm.listener.(*mockTokenListener)
+					log.Printf("FAILING WITH lastExecution [STARTED]:[LISTENER:%d]: %d", l.Id, lastExecution)
+				}
 				assert.Nil(t, tm.listener)
 			} else {
+				if tm.listener == nil {
+					log.Printf("FAILING WITH lastExecution[STOPPED]: %d", lastExecution)
+				}
 				assert.NotNil(t, tm.listener)
 				cancel, err := tokenManager.Start(listener)
 				assert.Nil(t, cancel)
@@ -606,7 +616,7 @@ func TestEntraidTokenManager_GetToken(t *testing.T) {
 		cancel, err := tokenManager.Start(listener)
 		assert.Error(t, err)
 		assert.Nil(t, cancel)
-		assert.NotNil(t, tm.listener)
+		assert.Nil(t, tm.listener)
 	})
 	t.Run("GetToken with expired manager", func(t *testing.T) {
 		t.Parallel()
@@ -1035,7 +1045,6 @@ func TestEntraidTokenManager_Streaming(t *testing.T) {
 		listener.On("OnTokenError", mock.Anything).Run(func(args mock.Arguments) {
 			err := args.Get(0)
 			assert.NotNil(t, err)
-			log.Printf("Found TOKEN Error: %v", err)
 		}).Return().Maybe()
 
 		cancel, err := tokenManager.Start(listener)
@@ -1109,7 +1118,7 @@ func TestEntraidTokenManager_Streaming(t *testing.T) {
 		assert.NotEqual(t, time.Duration(0), toRenewal)
 		assert.NotEqual(t, expiresIn, toRenewal)
 		assert.True(t, expiresIn > toRenewal)
-		<-time.After(toRenewal + 5*time.Millisecond)
+		<-time.After(toRenewal + 100*time.Millisecond)
 
 		idp.AssertNumberOfCalls(t, "RequestToken", 2)
 		listener.AssertNumberOfCalls(t, "OnTokenNext", 1)