add test for agent calling RefreshToken()

spikecurtis · spikecurtis · commit 80aa609f457f · 2025-09-02T11:18:36.000Z
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -2952,9 +2952,12 @@ func TestAgent_Reconnect(t *testing.T) {
 	defer closer.Close()
 
 	call1 := testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
+	require.Equal(t, client.GetNumRefreshTokenCalls(), 1)
 	close(call1.Resps) // hang up
 	// expect reconnect
 	testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
+	// Check that the agent refreshes the token when it reconnects.
+	require.Equal(t, client.GetNumRefreshTokenCalls(), 2)
 	closer.Close()
 }
 
diff --git a/agent/agenttest/client.go b/agent/agenttest/client.go
@@ -88,10 +88,11 @@ type Client struct {
 	fakeAgentAPI       *FakeAgentAPI
 	LastWorkspaceAgent func()
 
-	mu             sync.Mutex // Protects following.
-	logs           []agentsdk.Log
-	derpMapUpdates chan *tailcfg.DERPMap
-	derpMapOnce    sync.Once
+	mu                sync.Mutex // Protects following.
+	logs              []agentsdk.Log
+	derpMapUpdates    chan *tailcfg.DERPMap
+	derpMapOnce       sync.Once
+	refreshTokenCalls int
 }
 
 func (*Client) AsRequestOption() codersdk.RequestOption {
@@ -104,10 +105,19 @@ func (*Client) GetSessionToken() string {
 	return "agenttest-token"
 }
 
-func (*Client) RefreshToken(context.Context) error {
+func (c *Client) RefreshToken(context.Context) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.refreshTokenCalls++
 	return nil
 }
 
+func (c *Client) GetNumRefreshTokenCalls() int {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.refreshTokenCalls
+}
+
 func (*Client) RewriteDERPMap(*tailcfg.DERPMap) {}
 
 func (c *Client) Close() {
