Add ping test

Author: joecorall

main.go

@@ -35,13 +35,13 @@ func init() {
 	slog.SetDefault(handler)
 }
 
-func startPingRoutine(ctx context.Context, wg *sync.WaitGroup, c *config.Config) {
+func startPingRoutine(ctx context.Context, wg *sync.WaitGroup, c *config.Config, interval time.Duration) {
 	defer wg.Done()
 
-	ticker := time.NewTicker(30 * time.Second)
+	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
 
-	slog.Info("Starting ping routine to GCE instance")
+	slog.Info("Starting ping routine to GCE instance", "interval", interval)
 
 	for {
 		select {
@@ -88,18 +88,13 @@ func main() {
 		c.PowerOnCooldown = 30
 	}
 
-	// Create context for graceful shutdown
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-
-	// Set up signal handling for graceful shutdown
 	sigChan := make(chan os.Signal, 1)
 	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
-
-	// Start ping routine
 	var wg sync.WaitGroup
 	wg.Add(1)
-	go startPingRoutine(ctx, &wg, c)
+	go startPingRoutine(ctx, &wg, c, 30*time.Second)
 
 	http.HandleFunc("/healthcheck", func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusOK)
@@ -127,7 +122,6 @@ func main() {
 		p.ServeHTTP(w, r)
 	})
 
-	// Start HTTP server in a goroutine
 	server := &http.Server{Addr: ":8080"}
 	go func() {
 		slog.Info("Server listening on :8080")
@@ -136,22 +130,16 @@ func main() {
 		}
 	}()
 
-	// Wait for shutdown signal
 	<-sigChan
 	slog.Info("Received shutdown signal, gracefully shutting down...")
-
-	// Cancel context to stop ping routine
 	cancel()
 
-	// Shutdown HTTP server
 	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer shutdownCancel()
-
 	if err := server.Shutdown(shutdownCtx); err != nil {
 		slog.Error("Server shutdown error", "err", err)
 	}
 
-	// Wait for ping routine to finish
 	wg.Wait()
 	slog.Info("Shutdown complete")
 }

main_test.go

@@ -0,0 +1,128 @@
+package main
+
+import (
+	"context"
+	"log/slog"
+	"net/http"
+	"os"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/libops/ppb/pkg/config"
+	"github.com/libops/ppb/pkg/machine"
+)
+
+func TestStartPingRoutine_Integration(t *testing.T) {
+	// Track ping requests
+	var pingCount int
+	var pingURLs []string
+	var mu sync.Mutex
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/ping", func(w http.ResponseWriter, r *http.Request) {
+		mu.Lock()
+		defer mu.Unlock()
+		pingCount++
+		pingURLs = append(pingURLs, r.URL.Path)
+
+		w.WriteHeader(http.StatusOK)
+		_, err := w.Write([]byte("pong"))
+		if err != nil {
+			slog.Error("Unable to write ping response", "err", err)
+			os.Exit(1)
+		}
+	})
+	server := &http.Server{
+		Addr:    ":8808",
+		Handler: mux,
+	}
+	go func() {
+		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+			// Ignore "address already in use" errors during testing
+			if !strings.Contains(err.Error(), "address already in use") {
+				t.Errorf("Mock server error: %v", err)
+			}
+		}
+	}()
+
+	time.Sleep(1 * time.Second)
+
+	defer func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		err := server.Shutdown(ctx)
+		if err != nil {
+			slog.Error("Server shutdown failed", "err", err)
+			os.Exit(1)
+		}
+	}()
+
+	mockMachine := machine.NewGceMachine()
+	mockMachine.SetHostForTesting("127.0.0.1")
+
+	config := &config.Config{
+		Machine: mockMachine,
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 350*time.Millisecond)
+	defer cancel()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	go startPingRoutine(ctx, &wg, config, 100*time.Millisecond)
+	wg.Wait()
+
+	mu.Lock()
+	defer mu.Unlock()
+
+	if pingCount == 0 {
+		t.Error("Expected at least one ping request, got none")
+	}
+
+	// Verify all requests were to /ping endpoint
+	for _, url := range pingURLs {
+		if url != "/ping" {
+			t.Errorf("Expected ping to /ping endpoint, got %s", url)
+		}
+	}
+
+	// Should have made multiple pings (at least 2-3 in 350ms with 100ms interval)
+	if pingCount < 2 {
+		t.Errorf("Expected at least 2 ping requests, got %d", pingCount)
+	}
+}
+
+func TestStartPingRoutine_ContextCancellation(t *testing.T) {
+	mockMachine := machine.NewGceMachine()
+	mockMachine.SetHostForTesting("127.0.0.1")
+
+	config := &config.Config{
+		Machine: mockMachine,
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	routineFinished := make(chan bool, 1)
+
+	go func() {
+		startPingRoutine(ctx, &wg, config, 50*time.Millisecond)
+		routineFinished <- true
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+	cancel()
+
+	wg.Wait()
+
+	// Verify the routine actually finished
+	select {
+	case <-routineFinished:
+	case <-time.After(1 * time.Second):
+		t.Error("Ping routine did not finish within expected time after context cancellation")
+	}
+}