package healthz

import (
	"context"
	"fmt"
	"net/http"
	"testing"
	"time"
)

func TestHealthzReturns200WhileAlive(t *testing.T) {
	srv, err := New(0)
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}

	go func() {
		if err := srv.Start(); err != nil && err != http.ErrServerClosed {
			t.Errorf("Start returned unexpected error: %v", err)
		}
	}()

	// Give the server a moment to start listening.
	time.Sleep(50 * time.Millisecond)

	url := fmt.Sprintf("http://%s/healthz", srv.Addr())
	resp, err := http.Get(url)
	if err != nil {
		t.Fatalf("GET /healthz failed: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Errorf("/healthz status = %d, want %d", resp.StatusCode, http.StatusOK)
	}

	if err := srv.Stop(context.Background()); err != nil {
		t.Fatalf("Stop failed: %v", err)
	}
}

func TestHealthzReturns503WhileShuttingDown(t *testing.T) {
	srv, err := New(0)
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}

	go func() {
		if err := srv.Start(); err != nil && err != http.ErrServerClosed {
			t.Errorf("Start returned unexpected error: %v", err)
		}
	}()

	time.Sleep(50 * time.Millisecond)

	// Initiate shutdown but don't wait for it to finish.
	shutdownCtx, cancel := context.WithCancel(context.Background())
	go func() {
		_ = srv.Stop(shutdownCtx)
	}()

	// Give the shutdown flag time to flip.
	time.Sleep(50 * time.Millisecond)

	url := fmt.Sprintf("http://%s/healthz", srv.Addr())
	resp, err := http.Get(url)
	if err != nil {
		t.Fatalf("GET /healthz during shutdown failed: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Errorf("/healthz status during shutdown = %d, want %d", resp.StatusCode, http.StatusServiceUnavailable)
	}

	cancel()
}
