From c0aaedfc3bb332b3b4964063f28f0e7686473762 Mon Sep 17 00:00:00 2001 From: perennial Date: Thu, 19 Sep 2024 19:51:01 +1000 Subject: [PATCH] Add test suite for token_manager This commit introduces a new file token_manager_test.go containing unit tests for the token_manager package. The test suite covers various aspects, including initialisation, token selection methods, status updates, fallback mechanisms, and concurrent access. --- server/token_manager/token_manager_test.go | 242 +++++++++++++++++++++ 1 file changed, 242 insertions(+) create mode 100644 server/token_manager/token_manager_test.go diff --git a/server/token_manager/token_manager_test.go b/server/token_manager/token_manager_test.go new file mode 100644 index 0000000..961230c --- /dev/null +++ b/server/token_manager/token_manager_test.go @@ -0,0 +1,242 @@ +package token_manager + +import ( + "sync" + "testing" + "time" +) + +// TestNewTokenManager verifies that the NewTokenManager function correctly initializes +// a TokenManager with the provided configuration. +func TestNewTokenManager(t *testing.T) { + tokenValues := []string{"token1", "token2", "token3"} + maxRetries := 3 + baseTimeout := 5 * time.Second + maxBackoffTime := 30 * time.Second + loadBalancingMethod := "round-robin" + + tm := NewTokenManager(tokenValues, maxRetries, baseTimeout, maxBackoffTime, loadBalancingMethod) + + // Check if the number of tokens matches the input + if len(tm.tokens) != len(tokenValues) { + t.Errorf("Expected %d tokens, got %d", len(tokenValues), len(tm.tokens)) + } + + // Verify that all configuration parameters are correctly set + if tm.maxRetries != maxRetries { + t.Errorf("Expected maxRetries to be %d, got %d", maxRetries, tm.maxRetries) + } + + if tm.baseTimeout != baseTimeout { + t.Errorf("Expected baseTimeout to be %v, got %v", baseTimeout, tm.baseTimeout) + } + + if tm.maxBackoffTime != maxBackoffTime { + t.Errorf("Expected maxBackoffTime to be %v, got %v", maxBackoffTime, tm.maxBackoffTime) + } + + if tm.loadBalancingMethod != loadBalancingMethod { + t.Errorf("Expected loadBalancingMethod to be %s, got %s", loadBalancingMethod, tm.loadBalancingMethod) + } +} + +// TestGetTokenAllMethods tests all implemented load balancing methods to ensure +// they behave as expected when selecting tokens. +func TestGetTokenAllMethods(t *testing.T) { + tests := []struct { + name string + loadBalancingMethod string + expectedBehavior func(*testing.T, *TokenManager) + }{ + { + name: "Round Robin", + loadBalancingMethod: "round-robin", + expectedBehavior: func(t *testing.T, tm *TokenManager) { + // Test if tokens are returned in a cyclic order + for i := 0; i < len(tm.tokens)*2; i++ { + token := tm.GetToken() + expectedValue := tm.tokens[i%len(tm.tokens)].Value + if token.Value != expectedValue { + t.Errorf("Expected token value %s, got %s", expectedValue, token.Value) + } + } + }, + }, + { + name: "Random", + loadBalancingMethod: "random", + expectedBehavior: func(t *testing.T, tm *TokenManager) { + // Test if all tokens are used over multiple selections + usedTokens := make(map[string]bool) + for i := 0; i < len(tm.tokens)*10; i++ { + token := tm.GetToken() + usedTokens[token.Value] = true + } + if len(usedTokens) != len(tm.tokens) { + t.Errorf("Random selection did not use all available tokens") + } + }, + }, + { + name: "Least Recently Used", + loadBalancingMethod: "least-recently-used", + expectedBehavior: func(t *testing.T, tm *TokenManager) { + // Test if tokens are returned in order of least recent use + firstToken := tm.GetToken() + time.Sleep(10 * time.Millisecond) + secondToken := tm.GetToken() + time.Sleep(10 * time.Millisecond) + thirdToken := tm.GetToken() + + if firstToken.Value == secondToken.Value || firstToken.Value == thirdToken.Value || secondToken.Value == thirdToken.Value { + t.Errorf("Least-recently-used selection returned duplicate tokens") + } + }, + }, + } + + // Run tests for each load balancing method + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tm := NewTokenManager([]string{"token1", "token2", "token3"}, 3, 5*time.Second, 30*time.Second, tt.loadBalancingMethod) + tt.expectedBehavior(t, tm) + }) + } +} + +// TestMarkTokenStatus verifies that the MarkTokenStatus method correctly +// updates a token's status and handles failure counts. +func TestMarkTokenStatus(t *testing.T) { + tm := NewTokenManager([]string{"token1"}, 3, 5*time.Second, 30*time.Second, "round-robin") + token := tm.GetToken() + + // Test marking a token as TimedOut + tm.MarkTokenStatus(token, TimedOut) + if token.Status != TimedOut { + t.Errorf("Expected token status to be TimedOut, got %v", token.Status) + } + + if token.FailureCount != 1 { + t.Errorf("Expected failure count to be 1, got %d", token.FailureCount) + } + + // Test marking a token as Good (should reset failure count) + tm.MarkTokenStatus(token, Good) + if token.Status != Good { + t.Errorf("Expected token status to be Good, got %v", token.Status) + } + + if token.FailureCount != 0 { + t.Errorf("Expected failure count to be reset to 0, got %d", token.FailureCount) + } +} + +// TestResetAllTokens checks if the ResetAllTokens method correctly +// resets all tokens to their initial good state. +func TestResetAllTokens(t *testing.T) { + tm := NewTokenManager([]string{"token1", "token2"}, 3, 5*time.Second, 30*time.Second, "round-robin") + + // Mark all tokens as TimedOut + for _, token := range tm.tokens { + tm.MarkTokenStatus(token, TimedOut) + } + + // Reset all tokens + tm.ResetAllTokens() + + // Check if all tokens are reset to Good status with 0 failure count + for _, token := range tm.tokens { + if token.Status != Good { + t.Errorf("Expected all tokens to have Good status, got %v", token.Status) + } + if token.FailureCount != 0 { + t.Errorf("Expected all tokens to have FailureCount 0, got %d", token.FailureCount) + } + } +} + +// TestGetterMethods ensures that the getter methods return the correct values +// that were set during TokenManager initialization. +func TestGetterMethods(t *testing.T) { + maxRetries := 3 + baseTimeout := 5 * time.Second + maxBackoffTime := 30 * time.Second + tm := NewTokenManager([]string{"token1"}, maxRetries, baseTimeout, maxBackoffTime, "round-robin") + + if tm.GetMaxRetries() != maxRetries { + t.Errorf("Expected GetMaxRetries to return %d, got %d", maxRetries, tm.GetMaxRetries()) + } + + if tm.GetBaseTimeout() != baseTimeout { + t.Errorf("Expected GetBaseTimeout to return %v, got %v", baseTimeout, tm.GetBaseTimeout()) + } + + if tm.GetMaxBackoffTime() != maxBackoffTime { + t.Errorf("Expected GetMaxBackoffTime to return %v, got %v", maxBackoffTime, tm.GetMaxBackoffTime()) + } +} + +// TestGetFallbackToken verifies that when all tokens are timed out, +// the TokenManager correctly selects and resets a fallback token. +func TestGetFallbackToken(t *testing.T) { + tm := NewTokenManager([]string{"token1", "token2"}, 3, 5*time.Second, 30*time.Second, "round-robin") + + // Mark all tokens as timed out + for _, token := range tm.tokens { + tm.MarkTokenStatus(token, TimedOut) + token.TimeoutUntil = time.Now().Add(-1 * time.Second) // Set timeout in the past + } + + // Get a token, which should reset and return a previously timed-out token + token := tm.GetToken() + if token == nil { + t.Error("Expected a fallback token, got nil") + } + if token.Status != Good { + t.Errorf("Expected fallback token status to be Good, got %v", token.Status) + } +} + +// TestExponentialBackoff checks if the exponential backoff mechanism +// correctly increases the timeout duration for consecutive failures. +func TestExponentialBackoff(t *testing.T) { + tm := NewTokenManager([]string{"token1"}, 3, 1*time.Second, 8*time.Second, "round-robin") + token := tm.GetToken() + + expectedTimeouts := []time.Duration{1 * time.Second, 2 * time.Second, 4 * time.Second, 8 * time.Second, 8 * time.Second} + + for i, expected := range expectedTimeouts { + tm.MarkTokenStatus(token, TimedOut) + if token.TimeoutUntil.Sub(time.Now()).Round(time.Second) != expected { + t.Errorf("Iteration %d: Expected timeout duration %v, got %v", i, expected, token.TimeoutUntil.Sub(time.Now()).Round(time.Second)) + } + } +} + +// TestConcurrentAccess verifies that the TokenManager can handle +// concurrent access from multiple goroutines without race conditions. +func TestConcurrentAccess(t *testing.T) { + tm := NewTokenManager([]string{"token1", "token2", "token3"}, 3, 5*time.Second, 30*time.Second, "round-robin") + + var wg sync.WaitGroup + for i := 0; i < 100; i++ { + wg.Add(1) + go func() { + defer wg.Done() + token := tm.GetToken() + tm.MarkTokenStatus(token, Good) + }() + } + wg.Wait() +} + +// TestEmptyTokenList checks if the TokenManager correctly handles +// the case when initialized with an empty list of tokens. +func TestEmptyTokenList(t *testing.T) { + tm := NewTokenManager([]string{}, 3, 5*time.Second, 30*time.Second, "round-robin") + + token := tm.GetToken() + if token != nil { + t.Errorf("Expected nil token for empty token list, got %v", token) + } +}