homelab-frontend/internal/testsupport/stub.go

package testsupport

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"
)

// Stub response paths. testdata/config/harness.yaml points every upstream at the
// one stub server, so the response shape is selected by path, not by port.
const (
	PathJSON    = "/stub/json"
	PathSSE     = "/stub/sse"
	PathChunked = "/stub/chunked"
	PathSlow    = "/stub/slow"
)

// SSEGap is the pause between SSE events. It exists so a test can observe that
// chunks arrive incrementally rather than all at once at the end.
const SSEGap = 20 * time.Millisecond

// SlowDelay is how long PathSlow waits before writing anything.
const SlowDelay = 250 * time.Millisecond

// SSETokens are the tokens PathSSE emits, one event per token, followed by the
// [DONE] sentinel that OpenAI-shaped clients expect.
var SSETokens = []string{"Hello", " ", "world", "!"}

// ChunkedBodies are the pieces PathChunked writes, each flushed separately so
// the response goes out with Transfer-Encoding: chunked.
var ChunkedBodies = []string{"first\n", "second\n", "third\n"}

// newStubHandler builds the mux served by the harness. Every handler writes a
// deterministic body so tests can assert on exact bytes.
func newStubHandler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc(PathJSON, stubJSON)
	mux.HandleFunc(PathSSE, stubSSE)
	mux.HandleFunc(PathChunked, stubChunked)
	mux.HandleFunc(PathSlow, stubSlow)
	return mux
}

// stubJSON serves a fixed JSON body.
func stubJSON(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	// Encode failure here means the client hung up mid-write; the connection is
	// already gone, so there is nothing to report and no header left to change.
	_ = json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}

// stubSSE streams one event per token, flushing after each. It returns early
// when the client disconnects, so a mid-response disconnect test can assert on
// how many tokens the upstream actually managed to emit.
func stubSSE(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.WriteHeader(http.StatusOK)

	rc := http.NewResponseController(w)
	if err := rc.Flush(); err != nil {
		return
	}

	for _, tok := range SSETokens {
		select {
		case <-r.Context().Done():
			return
		case <-time.After(SSEGap):
		}

		if _, err := fmt.Fprintf(w, "data: %s\n\n", tok); err != nil {
			return
		}
		if err := rc.Flush(); err != nil {
			return
		}
	}

	if _, err := fmt.Fprint(w, "data: [DONE]\n\n"); err != nil {
		return
	}
	_ = rc.Flush()
}

// stubChunked writes several pieces with a flush between each, producing a
// chunked transfer with no Content-Length.
func stubChunked(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)

	rc := http.NewResponseController(w)
	for _, body := range ChunkedBodies {
		select {
		case <-r.Context().Done():
			return
		case <-time.After(SSEGap):
		}

		if _, err := fmt.Fprint(w, body); err != nil {
			return
		}
		if err := rc.Flush(); err != nil {
			return
		}
	}
}

// stubSlow waits SlowDelay before responding at all, for timeout tests.
func stubSlow(w http.ResponseWriter, r *http.Request) {
	select {
	case <-r.Context().Done():
		return
	case <-time.After(SlowDelay):
	}

	w.Header().Set("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	_, _ = fmt.Fprint(w, "slow\n")
}