Test that SSH keepalives increase the expiry time (#50503)

lib/inventory/controller_test.go

@@ -61,6 +61,14 @@ type fakeAuth struct {
 
 	lastInstance    types.Instance
 	lastRawInstance []byte
+
+	lastServerExpiry time.Time
+}
+
+func (a *fakeAuth) getLastServerExpiry() time.Time {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	return a.lastServerExpiry
 }
 
 func (a *fakeAuth) UpsertNode(_ context.Context, server types.Server) (*types.KeepAlive, error) {
@@ -76,6 +84,7 @@ func (a *fakeAuth) UpsertNode(_ context.Context, server types.Server) (*types.Ke
 		a.failUpserts--
 		return nil, trace.Errorf("upsert failed as test condition")
 	}
+	a.lastServerExpiry = server.Expiry()
 	return &types.KeepAlive{}, a.err
 }
 
@@ -88,6 +97,7 @@ func (a *fakeAuth) UpsertApplicationServer(_ context.Context, server types.AppSe
 		a.failUpserts--
 		return nil, trace.Errorf("upsert failed as test condition")
 	}
+	a.lastServerExpiry = server.Expiry()
 	return &types.KeepAlive{}, a.err
 }
 
@@ -104,6 +114,7 @@ func (a *fakeAuth) UpsertDatabaseServer(_ context.Context, server types.Database
 		a.failUpserts--
 		return nil, trace.Errorf("upsert failed as test condition")
 	}
+	a.lastServerExpiry = server.Expiry()
 	return &types.KeepAlive{}, a.err
 }
 
@@ -120,21 +131,23 @@ func (a *fakeAuth) UpsertKubernetesServer(_ context.Context, server types.KubeSe
 		a.failUpserts--
 		return nil, trace.Errorf("upsert failed as test condition")
 	}
+	a.lastServerExpiry = server.Expiry()
 	return &types.KeepAlive{}, a.err
 }
 
 func (a *fakeAuth) DeleteKubernetesServer(ctx context.Context, hostID, name string) error {
 	return nil
 }
 
-func (a *fakeAuth) KeepAliveServer(_ context.Context, _ types.KeepAlive) error {
+func (a *fakeAuth) KeepAliveServer(_ context.Context, ka types.KeepAlive) error {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 	a.keepalives++
 	if a.failKeepAlives > 0 {
 		a.failKeepAlives--
 		return trace.Errorf("keepalive failed as test condition")
 	}
+	a.lastServerExpiry = ka.Expires
 	return a.err
 }
 
@@ -222,6 +235,10 @@ func TestSSHServerBasics(t *testing.T) {
 		deny(sshUpsertErr, sshKeepAliveErr, handlerClose),
 	)
 
+	// we will check that the expiration time will grow after keepalives and new
+	// server announces
+	expiry := auth.getLastServerExpiry()
+
 	// set up to induce some failures, but not enough to cause the control
 	// stream to be closed.
 	auth.mu.Lock()
@@ -254,6 +271,9 @@ func TestSSHServerBasics(t *testing.T) {
 		deny(sshKeepAliveErr, sshUpsertErr, sshUpsertRetryOk, handlerClose),
 	)
 
+	oldExpiry, expiry := expiry, auth.getLastServerExpiry()
+	require.Greater(t, expiry, oldExpiry)
+
 	err = downstream.Send(ctx, proto.InventoryHeartbeat{
 		SSHServer: &types.ServerV2{
 			Metadata: types.Metadata{
@@ -292,6 +312,9 @@ func TestSSHServerBasics(t *testing.T) {
 		}
 	}()
 
+	oldExpiry, expiry = expiry, auth.getLastServerExpiry()
+	require.Greater(t, expiry, oldExpiry)
+
 	// limit time of ping call
 	pingCtx, cancel := context.WithTimeout(ctx, time.Second*10)
 	defer cancel()