diff --git a/vault/counters_test.go b/vault/counters_test.go index 164d58cda..3e9724e46 100644 --- a/vault/counters_test.go +++ b/vault/counters_test.go @@ -123,7 +123,7 @@ func TestRequestCounterSaveCurrent(t *testing.T) { } } -func testCountActiveTokens(t *testing.T, c *Core, root string, expectedServiceTokens int) { +func testCountActiveTokens(t *testing.T, c *Core, root string) int { t.Helper() rootCtx := namespace.RootContext(nil) @@ -136,15 +136,8 @@ func testCountActiveTokens(t *testing.T, c *Core, root string, expectedServiceTo t.Fatalf("bad: resp: %#v\n err: %v", resp, err) } - if diff := deep.Equal(resp.Data, map[string]interface{}{ - "counters": &ActiveTokens{ - ServiceTokens: TokenCounter{ - Total: expectedServiceTokens, - }, - }, - }); diff != nil { - t.Fatal(diff) - } + activeTokens := resp.Data["counters"].(*ActiveTokens) + return activeTokens.ServiceTokens.Total } func TestTokenStore_CountActiveTokens(t *testing.T) { @@ -152,7 +145,10 @@ func TestTokenStore_CountActiveTokens(t *testing.T) { rootCtx := namespace.RootContext(nil) // Count the root token - testCountActiveTokens(t, c, root, 1) + count := testCountActiveTokens(t, c, root) + if count != 1 { + t.Fatalf("expected %d tokens, not %d", 1, count) + } // Create some service tokens req := &logical.Request{ @@ -171,7 +167,10 @@ func TestTokenStore_CountActiveTokens(t *testing.T) { } tokens[i] = resp.Auth.ClientToken - testCountActiveTokens(t, c, root, i+2) + count = testCountActiveTokens(t, c, root) + if count != i+2 { + t.Fatalf("expected %d tokens, not %d", i+2, count) + } } // Revoke the service tokens @@ -183,9 +182,20 @@ func TestTokenStore_CountActiveTokens(t *testing.T) { if err != nil || (resp != nil && resp.IsError()) { t.Fatalf("bad: resp: %#v\n err: %v", resp, err) } - - testCountActiveTokens(t, c, root, 10-i) } + + // We should now have only 1 token (the root token). However, because + // token deletion works by setting the TTL of the token to 0 and waiting + // for it to get cleaned up by the expiration manager, occasionally we will + // have to wait briefly for all the tokens to actually get deleted. + for i := 0; i < 10; i++ { + count = testCountActiveTokens(t, c, root) + if count == 1 { + return + } + time.Sleep(time.Second) + } + t.Fatalf("expected %d tokens, not %d", 1, count) } func testCountActiveEntities(t *testing.T, c *Core, root string, expectedEntities int) {