chore: speed up tests (#241)

Author: jajeffries

Makefile

@@ -47,6 +47,19 @@ agent_bin:
 test:
 	@go test -race ./...
 
+.PHONY: test-timed
+test-timed:
+	@echo "Running tests with timing measurement..."
+	@echo "=========================================="
+	@start=$$(date +%s); \
+	go test -count=1 ./...; \
+	exit_code=$$?; \
+	end=$$(date +%s); \
+	duration=$$((end - start)); \
+	echo "=========================================="; \
+	echo "Test execution completed in $$duration seconds"; \
+	exit $$exit_code
+
 .PHONY: test-coverage
 test-coverage:
 	@mkdir -p .coverage

agent/configmgr/fleet.go

@@ -21,7 +21,7 @@ var _ Manager = (*fleetConfigManager)(nil)
 
 type fleetConfigManager struct {
 	logger            *slog.Logger
-	connection        *fleet.MQTTConnection
+	connection        fleet.MQTTConnector
 	authTokenManager  *fleet.AuthTokenManager
 	resetChan         chan struct{}
 	reconnectChan     chan struct{}
@@ -51,6 +51,21 @@ func newFleetConfigManager(logger *slog.Logger, pMgr policymgr.PolicyManager, ba
 	}
 }
 
+// newFleetConfigManagerWithConnection creates a fleetConfigManager with a custom connection (for testing)
+func newFleetConfigManagerWithConnection(logger *slog.Logger, pMgr policymgr.PolicyManager, backendState backend.StateRetriever, conn fleet.MQTTConnector) *fleetConfigManager {
+	resetChan := make(chan struct{}, 1)
+	reconnectChan := make(chan struct{}, 1)
+	return &fleetConfigManager{
+		logger:           logger,
+		connection:       conn, // Use provided connection instead of creating new one
+		authTokenManager: fleet.NewAuthTokenManager(logger),
+		resetChan:        resetChan,
+		reconnectChan:    reconnectChan,
+		backendState:     backendState,
+		policyManager:    pMgr,
+	}
+}
+
 func (fleetManager *fleetConfigManager) Start(cfg config.Config, backends map[string]backend.Backend) error {
 	ctx := context.Background()
 

agent/configmgr/fleet/auth_test.go

@@ -299,8 +299,7 @@ func TestAuthTokenManager_IsTokenExpired_ExpiredCases(t *testing.T) {
 	_, err := authTokenManager.GetToken(ctx, serverExpired.URL, true, 60*time.Second, "test_client_id", "test_client_secret")
 	require.NoError(t, err)
 
-	// Wait a moment to ensure time has passed
-	time.Sleep(2 * time.Second)
+	// Token is already expired (pastExpiry is 1 hour ago), so no wait needed
 	assert.True(t, authTokenManager.IsTokenExpired())
 }
 

agent/configmgr/fleet/connection.go

@@ -66,6 +66,14 @@ type ConnectionDetails struct {
 	Zone     string
 }
 
+// MQTTConnector defines the interface for MQTT connection operations
+type MQTTConnector interface {
+	Connect(ctx context.Context, details ConnectionDetails, backends map[string]backend.Backend, labels map[string]string, configFile string) error
+	Disconnect(ctx context.Context, heartbeatTopic string) error
+	Reconnect(ctx context.Context, details ConnectionDetails, backends map[string]backend.Backend, labels map[string]string, configFile string, timeout time.Duration) error
+	AddOnReadyHook(fn func(cm *autopaho.ConnectionManager, topics TokenResponseTopics))
+}
+
 // Connect connects to the MQTT broker
 func (connection *MQTTConnection) Connect(ctx context.Context, details ConnectionDetails, backends map[string]backend.Backend, labels map[string]string, configFile string) error {
 	// Parse the ORB URL

agent/configmgr/fleet/connection_test.go

@@ -101,8 +101,8 @@ func TestFleetConfigManager_Connect_ValidURL(t *testing.T) {
 	// since we don't have a real MQTT server
 	backends := make(map[string]backend.Backend)
 	trt2 := TokenResponseTopics{Inbox: "test/topic"}
-	// Timeout after 3 seconds
-	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	// Timeout after 100ms for faster test execution (connection will fail quickly)
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 	defer cancel()
 	err := connection.Connect(ctx,
 		ConnectionDetails{MQTTURL: "mqtt://localhost:1883", Token: "test_token", AgentID: "test-agent-id", Topics: trt2, ClientID: "test-agent-id", Zone: "test-zone"},

agent/configmgr/fleet/mocks.go

@@ -7,6 +7,7 @@ import (
 	"log/slog"
 	"time"
 
+	"github.com/eclipse/paho.golang/autopaho"
 	"github.com/stretchr/testify/mock"
 
 	"github.com/netboxlabs/orb-agent/agent/backend"
@@ -86,3 +87,40 @@ func (m *mockBackend) RemovePolicy(data policies.PolicyData) error {
 	args := m.Called(data)
 	return args.Error(0)
 }
+
+// MockMQTTConnection is a mock implementation of MQTTConnector for testing
+type MockMQTTConnection struct {
+	ConnectError    error
+	DisconnectError error
+	ReconnectError  error
+	ConnectCalled   bool
+	hooks           []func(cm *autopaho.ConnectionManager, topics TokenResponseTopics)
+}
+
+// Connect connects to the MQTT broker
+func (m *MockMQTTConnection) Connect(_ context.Context, _ ConnectionDetails, _ map[string]backend.Backend, _ map[string]string, _ string) error {
+	m.ConnectCalled = true
+	return m.ConnectError
+}
+
+// Disconnect disconnects from the MQTT broker
+func (m *MockMQTTConnection) Disconnect(_ context.Context, _ string) error {
+	return m.DisconnectError
+}
+
+// Reconnect reconnects to the MQTT broker
+func (m *MockMQTTConnection) Reconnect(_ context.Context, _ ConnectionDetails, _ map[string]backend.Backend, _ map[string]string, _ string, _ time.Duration) error {
+	return m.ReconnectError
+}
+
+// AddOnReadyHook registers a hook function to be called when the MQTT connection is ready.
+func (m *MockMQTTConnection) AddOnReadyHook(fn func(cm *autopaho.ConnectionManager, topics TokenResponseTopics)) {
+	m.hooks = append(m.hooks, fn)
+}
+
+// TriggerOnReadyHook triggers all registered onReady hooks (for testing)
+func (m *MockMQTTConnection) TriggerOnReadyHook(cm *autopaho.ConnectionManager, topics TokenResponseTopics) {
+	for _, hook := range m.hooks {
+		hook(cm, topics)
+	}
+}

agent/configmgr/fleet_test.go

@@ -3,6 +3,7 @@ package configmgr
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"log/slog"
 	"net/http"
 	"net/http/httptest"
@@ -107,15 +108,32 @@ func TestFleetConfigManager_Start_ConnectError(t *testing.T) {
 	// Arrange
 	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
 	mockPMgr := &mockPolicyManagerForFleet{}
-	fleetManager := newFleetConfigManager(logger, mockPMgr, &mockBackendState{})
+
+	// Create mock MQTT connection that returns a connection error immediately
+	mockConn := &fleet.MockMQTTConnection{
+		ConnectError: fmt.Errorf("mqtt connection failed: invalid URL"),
+	}
+	fleetManager := newFleetConfigManagerWithConnection(logger, mockPMgr, &mockBackendState{}, mockConn)
 
 	// Create mock HTTP server for token endpoint
+	// Use a valid JWT token so parsing succeeds and we reach the connection step
+	validJWT := fleet.RawJWTWithClaims(map[string]any{
+		"orb:org_id":   "test-org",
+		"orb:zone":     "default",
+		"orb:agent_id": "test-agent",
+		"client_id":    "test-client",
+		"iat":          1516239022,
+		"ext": map[string]any{
+			"orb:mqtt_url": "mqtt://test.example.com:1883",
+		},
+	})
 	server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
 		response := fleet.TokenResponse{
-			AccessToken: "test_access_token",
-			MQTTURL:     "://invalid-mqtt-url", // Invalid MQTT URL
+			AccessToken: validJWT,
+			MQTTURL:     "mqtt://test.example.com:1883",
 			ExpiresIn:   3600,
 		}
+		w.Header().Set("Content-Type", "application/json")
 		_ = json.NewEncoder(w).Encode(response)
 	}))
 	defer server.Close()
@@ -126,6 +144,7 @@ func TestFleetConfigManager_Start_ConnectError(t *testing.T) {
 				Sources: config.Sources{
 					Fleet: config.FleetManager{
 						TokenURL:     server.URL,
+						SkipTLS:      true, // Skip TLS verification for test server
 						ClientID:     "test_client",
 						ClientSecret: "test_secret",
 					},
@@ -141,6 +160,8 @@ func TestFleetConfigManager_Start_ConnectError(t *testing.T) {
 
 	// Assert
 	assert.Error(t, err)
+	// Verify the error is from the mock connection (no 30s wait)
+	assert.Contains(t, err.Error(), "mqtt connection failed")
 }
 
 func TestFleetConfigManager_GetContext(t *testing.T) {
@@ -182,7 +203,12 @@ func TestFleetConfigManager_Start_WithJWTTopicGeneration(t *testing.T) {
 
 	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug}))
 	mockPMgr := &mockPolicyManagerForFleet{}
-	fleetManager := newFleetConfigManager(logger, mockPMgr, &mockBackendState{})
+
+	// Create mock MQTT connection that returns a connection error
+	mockConn := &fleet.MockMQTTConnection{
+		ConnectError: fmt.Errorf("mqtt connection failed: connection refused"),
+	}
+	fleetManager := newFleetConfigManagerWithConnection(logger, mockPMgr, &mockBackendState{}, mockConn)
 
 	// Create mock HTTP server that returns a JWT token
 	server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -216,22 +242,19 @@ func TestFleetConfigManager_Start_WithJWTTopicGeneration(t *testing.T) {
 	// Mock backends
 	backends := make(map[string]backend.Backend)
 
-	// Since we can't easily mock the MQTT connection in this test,
-	// we expect the Start method to fail at the MQTT connection step,
-	// but succeed in generating topics from JWT claims
+	// The Start method should succeed in generating topics from JWT claims,
+	// but fail at the MQTT connection step (which is mocked to return an error)
 	err := fleetManager.Start(cfg, backends)
 
-	// We expect an error because we can't actually connect to MQTT,
+	// We expect an error because the mock connection returns an error,
 	// but we want to verify that topic generation succeeded
 	// (The error should be related to MQTT connection, not JWT parsing)
 	require.Error(t, err)
-	// The error could be connection-related (mqtt, tcp, timeout, etc.)
+	// The error should be the mocked connection error
 	errorMsg := strings.ToLower(err.Error())
 	assert.True(t,
 		strings.Contains(errorMsg, "mqtt") ||
-			strings.Contains(errorMsg, "connection") ||
-			strings.Contains(errorMsg, "timeout") ||
-			strings.Contains(errorMsg, "deadline"),
+			strings.Contains(errorMsg, "connection"),
 		"Expected connection-related error, got: %s", err.Error())
 }
 
@@ -386,7 +409,7 @@ func TestFleetConfigManager_MonitorTokenExpiry_Configuration(t *testing.T) {
 	select {
 	case <-done:
 		// Monitor stopped successfully
-	case <-time.After(1 * time.Second):
+	case <-time.After(100 * time.Millisecond):
 		t.Fatal("monitor did not stop within timeout")
 	}
 }
@@ -434,7 +457,7 @@ func TestFleetConfigManager_MonitorTokenExpiry_CustomConfiguration(t *testing.T)
 	select {
 	case <-done:
 		// Monitor stopped successfully
-	case <-time.After(1 * time.Second):
+	case <-time.After(100 * time.Millisecond):
 		t.Fatal("monitor did not stop within timeout")
 	}
 }
@@ -467,7 +490,7 @@ func TestFleetConfigManager_Stop_CancelsMonitor(t *testing.T) {
 	select {
 	case <-fleetManager.monitorCtx.Done():
 		// Good, context is cancelled
-	case <-time.After(100 * time.Millisecond):
+	case <-time.After(50 * time.Millisecond):
 		t.Fatal("monitor context should be cancelled after Stop()")
 	}
 }
@@ -572,10 +595,11 @@ func TestFleetConfigManager_MonitorTokenExpiry_DetectsExpiredToken(t *testing.T)
 	}()
 
 	// Wait for monitor to detect expired token
+	// Allow at least one ticker interval (100ms) plus buffer
 	select {
 	case <-done:
 		// Monitor detected expired token
-	case <-time.After(500 * time.Millisecond):
+	case <-time.After(200 * time.Millisecond):
 		t.Fatal("monitor did not detect expired token within timeout")
 	}
 
@@ -670,10 +694,11 @@ func TestFleetConfigManager_MonitorTokenExpiry_DetectsExpiringSoonToken(t *testi
 	}()
 
 	// Wait for monitor to detect expiring token
+	// Allow at least one ticker interval (100ms) plus buffer
 	select {
 	case <-done:
 		// Monitor detected expiring token
-	case <-time.After(500 * time.Millisecond):
+	case <-time.After(200 * time.Millisecond):
 		t.Fatal("monitor did not detect expiring token within timeout")
 	}
 

agent/secretsmgr/vault_auth_test.go

@@ -51,9 +51,8 @@ func TestNewAuthentication(t *testing.T) {
 }
 
 func TestTokenAuth_Authenticate(t *testing.T) {
-	// Create test vault server
-	cluster, client := createTestVault(t)
-	defer cluster.Cleanup()
+	// Use shared test vault server
+	_, client := createTestVault(t)
 
 	ctx := context.Background()