From d3bb5d60b202dc0a5a85d31b541ccc7878b97169 Mon Sep 17 00:00:00 2001
From: Mike Jarmy <mjarmy@gmail.com>
Date: Wed, 30 Oct 2019 08:03:03 -0400
Subject: [PATCH] Fix the token counter test so it doesn't miscount the number
 of tokens. (#7760)

---
 vault/counters_test.go | 38 ++++++++++++++++++++++++--------------
 1 file changed, 24 insertions(+), 14 deletions(-)

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) {