ha-n95-local-control/cmd/n95bridge/shutdown_test.go

package main

import (
	"bytes"
	"context"
	"fmt"
	"net"
	"net/http"
	"strings"
	"sync"
	"testing"
	"time"
)

func getFreePort(t *testing.T) int {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("get free port: %v", err)
	}
	defer ln.Close()
	return ln.Addr().(*net.TCPAddr).Port
}

func TestGracefulStopPublishesOffline(t *testing.T) {
	pLookup := getFreePort(t)
	pFirmware := getFreePort(t)
	pXMPP := getFreePort(t)
	pHealth := getFreePort(t)

	env := []string{
		"ADVERTISE_IP=192.0.2.10",
		"MQTT_HOST=mqtt.example.invalid",
		fmt.Sprintf("PORT_LOOKUP=%d", pLookup),
		fmt.Sprintf("PORT_FIRMWARE=%d", pFirmware),
		fmt.Sprintf("PORT_XMPP=%d", pXMPP),
		fmt.Sprintf("HEALTH_PORT=%d", pHealth),
	}

	var mu sync.Mutex
	var order []string
	var bridgeDeadline time.Time
	var bridgeHasDeadline bool

	fakeShutdownBridge := func(ctx context.Context) error {
		mu.Lock()
		defer mu.Unlock()
		dl, ok := ctx.Deadline()
		bridgeHasDeadline = ok
		bridgeDeadline = dl
		order = append(order, "bridge_published_retained_offline")
		return nil
	}

	fakeShutdownXMPP := func(ctx context.Context) error {
		mu.Lock()
		defer mu.Unlock()
		order = append(order, "xmpp_stream_closed")
		return nil
	}

	shutdownBridge = fakeShutdownBridge
	shutdownXMPP = fakeShutdownXMPP
	t.Cleanup(func() {
		shutdownBridge = nil
		shutdownXMPP = nil
	})

	ctx, cancel := context.WithCancel(context.Background())
	runDone := make(chan error, 1)
	go func() {
		runDone <- run(ctx, env)
	}()

	// Wait for health endpoint to become ready.
	healthURL := fmt.Sprintf("http://127.0.0.1:%d/healthz", pHealth)
	ready := false
	for i := 0; i < 50; i++ {
		time.Sleep(20 * time.Millisecond)
		resp, err := http.Get(healthURL)
		if err == nil {
			resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				ready = true
				break
			}
		}
	}
	if !ready {
		t.Fatal("servers did not become ready")
	}

	// Trigger graceful stop (simulating SIGTERM via root context cancel).
	cancel()

	select {
	case err := <-runDone:
		if err != nil {
			t.Fatalf("run returned error on graceful shutdown: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("run did not exit within 5 second shutdown budget")
	}

	mu.Lock()
	defer mu.Unlock()

	if len(order) != 2 {
		t.Fatalf("expected 2 shutdown actions, got %d: %v", len(order), order)
	}
	if order[0] != "bridge_published_retained_offline" {
		t.Errorf("expected step 1 to be bridge publishing offline, got: %s", order[0])
	}
	if order[1] != "xmpp_stream_closed" {
		t.Errorf("expected step 2 to be xmpp stream closed, got: %s", order[1])
	}

	if !bridgeHasDeadline {
		t.Error("shutdownBridge context did not have a timeout deadline")
	} else {
		remaining := time.Until(bridgeDeadline)
		if remaining > 5*time.Second {
			t.Errorf("shutdown timeout deadline > 5s: %v", remaining)
		}
	}
}

type safeBuffer struct {
	mu  sync.Mutex
	buf bytes.Buffer
}

func (b *safeBuffer) Write(p []byte) (n int, err error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buf.Write(p)
}

func (b *safeBuffer) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buf.String()
}

func TestShutdownBudgetExpiredLogsAndExitsZero(t *testing.T) {
	origTimeout := shutdownTimeout
	shutdownTimeout = 50 * time.Millisecond
	defer func() { shutdownTimeout = origTimeout }()

	logBuf := &safeBuffer{}
	origLogOutput := logOutput
	logOutput = logBuf
	defer func() { logOutput = origLogOutput }()

	pLookup := getFreePort(t)
	pFirmware := getFreePort(t)
	pXMPP := getFreePort(t)
	pHealth := getFreePort(t)

	env := []string{
		"ADVERTISE_IP=192.0.2.10",
		"MQTT_HOST=mqtt.example.invalid",
		fmt.Sprintf("PORT_LOOKUP=%d", pLookup),
		fmt.Sprintf("PORT_FIRMWARE=%d", pFirmware),
		fmt.Sprintf("PORT_XMPP=%d", pXMPP),
		fmt.Sprintf("HEALTH_PORT=%d", pHealth),
	}

	// Slow bridge shutdown that exceeds the budget.
	shutdownBridge = func(ctx context.Context) error {
		<-ctx.Done() // wait for timeout
		return ctx.Err()
	}
	shutdownXMPP = func(ctx context.Context) error {
		return nil
	}
	t.Cleanup(func() {
		shutdownBridge = nil
		shutdownXMPP = nil
	})

	ctx, cancel := context.WithCancel(context.Background())
	runDone := make(chan error, 1)
	go func() {
		runDone <- run(ctx, env)
	}()

	// Wait for health endpoint.
	healthURL := fmt.Sprintf("http://127.0.0.1:%d/healthz", pHealth)
	for i := 0; i < 50; i++ {
		time.Sleep(20 * time.Millisecond)
		resp, err := http.Get(healthURL)
		if err == nil {
			resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				break
			}
		}
	}

	cancel()

	select {
	case err := <-runDone:
		if err != nil {
			t.Fatalf("expected exit 0 (nil error) on budget expiration, got: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("run did not return within expected window")
	}

	out := logBuf.String()
	if !strings.Contains(out, "shutdown budget expired") {
		t.Fatalf("expected log 'shutdown budget expired', got: %s", out)
	}
}