http-latency-prober/go/run_test.go

package main

import (
	"bytes"
	"encoding/json"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// ── helpers ───────────────────────────────────────────────────────────────────

// statusSrv starts a server that always replies with the given HTTP status.
func statusSrv(code int) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(code)
	}))
}

// blockingSrv starts a server whose handler blocks until the client disconnects.
// Using r.Context().Done() means the handler exits cleanly when the client drops,
// so srv.Close() never hangs.
func blockingSrv() *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		<-r.Context().Done()
	}))
}

// refusedURL returns an http:// URL on a port where nothing is listening.
// It binds a listener to get a free port, closes it immediately, then hands
// back that address — so any connect attempt is immediately refused.
func refusedURL() string {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic("refusedURL: " + err.Error())
	}
	addr := l.Addr().String()
	l.Close()
	return "http://" + addr
}

// invoke calls run() and returns stdout, stderr, and the exit code.
func invoke(args ...string) (stdout, stderr string, code int) {
	var outBuf, errBuf bytes.Buffer
	code = run(args, &outBuf, &errBuf)
	return outBuf.String(), errBuf.String(), code
}

// ── matrix ────────────────────────────────────────────────────────────────────

func TestRunMatrix(t *testing.T) {
	ok200 := statusSrv(200)
	defer ok200.Close()

	ok500 := statusSrv(500)
	defer ok500.Close()

	ok404 := statusSrv(404)
	defer ok404.Close()

	tlsSrv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(200)
	}))
	defer tlsSrv.Close()

	hangSrv := blockingSrv()
	defer hangSrv.Close()

	refused := refusedURL()

	cases := []struct {
		name     string
		args     []string
		wantCode int
		wantOut  []string // substrings required in stdout
		wantErr  []string // substrings required in stderr
	}{
		{
			// httptest uses 127.0.0.1 — Go skips DNS for loopback, so no DNS row.
			name:     "success 200",
			args:     []string{ok200.URL},
			wantCode: exitOK,
			wantOut:  []string{"(200)", "TCP connect", "Total"},
		},
		{
			name:     "HTTP 500 without --fail",
			args:     []string{ok500.URL},
			wantCode: exitOK,
			wantOut:  []string{"(500)", "Total"},
		},
		{
			name:     "HTTP 404 with --fail",
			args:     []string{"--fail", ok404.URL},
			wantCode: exitHTTP,
			wantOut:  []string{"404 ✗"},
		},
		{
			name:     "DNS failure",
			args:     []string{"https://this.will.never.resolve.invalid"},
			wantCode: exitDNS,
			wantOut:  []string{"✗ dns:"},
		},
		{
			name:     "connection refused",
			args:     []string{refused},
			wantCode: exitConnect,
			wantOut:  []string{"✗ connect:"},
		},
		{
			name:     "timeout",
			args:     []string{"--timeout", "200ms", hangSrv.URL},
			wantCode: exitTimeout,
			wantOut:  []string{"✗ timeout:"},
		},
		{
			name:     "TLS failure (self-signed cert rejected by default client)",
			args:     []string{tlsSrv.URL},
			wantCode: exitTLS,
			wantOut:  []string{"✗ tls:"},
		},
		{
			name:     "multiple URLs — highest exit code wins",
			args:     []string{ok200.URL, "https://no.such.host.for.test.invalid"},
			wantCode: exitDNS, // 2 > 0
			wantOut:  []string{"(200)", "✗ dns:"},
		},
		{
			name:     "sampling -n 3 all success",
			args:     []string{"-n", "3", ok200.URL},
			wantCode: exitOK,
			wantOut:  []string{"3 samples", "min", "avg", "max"},
		},
		{
			name:     "no args → usage",
			args:     []string{},
			wantCode: exitUsage,
			wantErr:  []string{"Usage:"},
		},
		{
			name:     "-h → usage exit 0",
			args:     []string{"-h"},
			wantCode: exitOK,
			wantErr:  []string{"Usage:"},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			stdout, stderr, code := invoke(tc.args...)

			if code != tc.wantCode {
				t.Errorf("exit code = %d, want %d\nstdout:\n%s\nstderr:\n%s",
					code, tc.wantCode, stdout, stderr)
			}
			for _, s := range tc.wantOut {
				if !strings.Contains(stdout, s) {
					t.Errorf("stdout missing %q\nstdout:\n%s", s, stdout)
				}
			}
			for _, s := range tc.wantErr {
				if !strings.Contains(stderr, s) {
					t.Errorf("stderr missing %q\nstderr:\n%s", s, stderr)
				}
			}
		})
	}
}

// ── JSON-specific assertions ──────────────────────────────────────────────────

func TestJSONSuccess(t *testing.T) {
	srv := statusSrv(200)
	defer srv.Close()

	stdout, _, code := invoke("--json", srv.URL)
	if code != exitOK {
		t.Fatalf("exit code = %d, want 0\nstdout: %s", code, stdout)
	}

	var results []struct {
		URL       string         `json:"url"`
		Status    int            `json:"status"`
		Succeeded int            `json:"succeeded"`
		Failed    int            `json:"failed"`
		Phases    map[string]any `json:"phases"`
		Errors    []any          `json:"errors"`
	}
	if err := json.Unmarshal([]byte(stdout), &results); err != nil {
		t.Fatalf("invalid JSON: %v\nstdout: %s", err, stdout)
	}
	if len(results) != 1 {
		t.Fatalf("expected 1 result, got %d", len(results))
	}
	r := results[0]
	if r.Status != 200 {
		t.Errorf("status = %d, want 200", r.Status)
	}
	if r.Succeeded != 1 {
		t.Errorf("succeeded = %d, want 1", r.Succeeded)
	}
	if r.Failed != 0 {
		t.Errorf("failed = %d, want 0", r.Failed)
	}
	if r.Phases["total"] == nil {
		t.Error("phases.total missing")
	}
	if len(r.Errors) > 0 {
		t.Errorf("unexpected errors: %v", r.Errors)
	}
}

func TestJSONDNSFailure(t *testing.T) {
	stdout, _, code := invoke("--json", "https://no.such.host.json.test.invalid")
	if code != exitDNS {
		t.Fatalf("exit code = %d, want %d\nstdout: %s", code, exitDNS, stdout)
	}

	var results []struct {
		Succeeded int `json:"succeeded"`
		Failed    int `json:"failed"`
		Errors    []struct {
			Phase   string `json:"phase"`
			Count   int    `json:"count"`
			Message string `json:"message"`
		} `json:"errors"`
	}
	if err := json.Unmarshal([]byte(stdout), &results); err != nil {
		t.Fatalf("invalid JSON: %v\nstdout: %s", err, stdout)
	}
	if len(results) != 1 {
		t.Fatalf("expected 1 result, got %d", len(results))
	}
	r := results[0]
	if r.Succeeded != 0 {
		t.Errorf("succeeded = %d, want 0", r.Succeeded)
	}
	if r.Failed != 1 {
		t.Errorf("failed = %d, want 1", r.Failed)
	}
	if len(r.Errors) == 0 {
		t.Fatal("errors array is empty")
	}
	if r.Errors[0].Phase != "dns" {
		t.Errorf("error phase = %q, want \"dns\"", r.Errors[0].Phase)
	}
	if r.Errors[0].Count != 1 {
		t.Errorf("error count = %d, want 1", r.Errors[0].Count)
	}
}

// ── concurrency tests ─────────────────────────────────────────────────────────

// TestRunConcurrentOrder verifies that with -c > 1 the URL blocks are printed
// in the original input order and exit codes are accumulated correctly.
func TestRunConcurrentOrder(t *testing.T) {
	a := statusSrv(200)
	defer a.Close()
	b := statusSrv(200)
	defer b.Close()
	c := statusSrv(200)
	defer c.Close()

	// Run with explicit high concurrency — all three URLs measured in parallel.
	stdout, _, code := invoke("-c", "3", a.URL, b.URL, c.URL)
	if code != exitOK {
		t.Fatalf("exit code = %d, want 0\nstdout:\n%s", code, stdout)
	}

	// Each URL block must appear and in the correct order.
	posA := strings.Index(stdout, a.URL)
	posB := strings.Index(stdout, b.URL)
	posC := strings.Index(stdout, c.URL)
	if posA < 0 || posB < 0 || posC < 0 {
		t.Fatalf("one or more URLs missing from stdout:\n%s", stdout)
	}
	if !(posA < posB && posB < posC) {
		t.Errorf("URLs out of order: posA=%d posB=%d posC=%d\nstdout:\n%s",
			posA, posB, posC, stdout)
	}
}

// TestConcurrentWorstCode confirms that worstCode aggregation is correct when
// both successes and failures run concurrently.
func TestConcurrentWorstCode(t *testing.T) {
	good := statusSrv(200)
	defer good.Close()

	// DNS failure mixed with a successful URL — highest code (2) must win.
	stdout, _, code := invoke("-c", "2",
		good.URL,
		"https://totally.bogus.domain.for.concurrency.test.invalid",
	)
	if code != exitDNS {
		t.Errorf("exit code = %d, want %d (DNS)\nstdout:\n%s", code, exitDNS, stdout)
	}
	// Both URLs must appear in output (good one before the failing one).
	if !strings.Contains(stdout, good.URL) {
		t.Errorf("stdout missing good URL\nstdout:\n%s", stdout)
	}
	if !strings.Contains(stdout, "✗ dns:") {
		t.Errorf("stdout missing DNS failure marker\nstdout:\n%s", stdout)
	}
}

// TestConcurrentJSONOrder verifies that JSON entries preserve input order under
// parallel execution.
func TestConcurrentJSONOrder(t *testing.T) {
	a := statusSrv(200)
	defer a.Close()
	b := statusSrv(200)
	defer b.Close()

	stdout, _, code := invoke("--json", "-c", "2", a.URL, b.URL)
	if code != exitOK {
		t.Fatalf("exit code = %d, want 0\nstdout:\n%s", code, stdout)
	}

	var results []struct {
		URL string `json:"url"`
	}
	if err := json.Unmarshal([]byte(stdout), &results); err != nil {
		t.Fatalf("invalid JSON: %v\nstdout: %s", err, stdout)
	}
	if len(results) != 2 {
		t.Fatalf("expected 2 results, got %d", len(results))
	}
	if results[0].URL != a.URL {
		t.Errorf("results[0].url = %q, want %q", results[0].URL, a.URL)
	}
	if results[1].URL != b.URL {
		t.Errorf("results[1].url = %q, want %q", results[1].URL, b.URL)
	}
}

func TestJSONSampling(t *testing.T) {
	srv := statusSrv(200)
	defer srv.Close()

	stdout, _, code := invoke("--json", "-n", "3", srv.URL)
	if code != exitOK {
		t.Fatalf("exit code = %d, want 0\nstdout: %s", code, stdout)
	}

	var results []struct {
		Succeeded int `json:"succeeded"`
		Phases    map[string]struct {
			MinMS float64 `json:"min_ms"`
			AvgMS float64 `json:"avg_ms"`
			MaxMS float64 `json:"max_ms"`
		} `json:"phases"`
	}
	if err := json.Unmarshal([]byte(stdout), &results); err != nil {
		t.Fatalf("invalid JSON: %v\nstdout: %s", err, stdout)
	}
	r := results[0]
	if r.Succeeded != 3 {
		t.Errorf("succeeded = %d, want 3", r.Succeeded)
	}
	total, ok := r.Phases["total"]
	if !ok {
		t.Fatal("phases.total missing")
	}
	if total.MinMS <= 0 {
		t.Errorf("total.min_ms = %f, want > 0", total.MinMS)
	}
	if total.MaxMS < total.MinMS {
		t.Errorf("total.max_ms (%f) < total.min_ms (%f)", total.MaxMS, total.MinMS)
	}
}