cookie ready

internal/render/engine_test.go

@@ -0,0 +1,56 @@
+package render
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"testing"
+
+	"github.com/a-h/templ"
+)
+
+type writeCountingResponseWriter struct {
+	header http.Header
+	writes int
+	body   []byte
+}
+
+func (w *writeCountingResponseWriter) Header() http.Header {
+	if w.header == nil {
+		w.header = make(http.Header)
+	}
+	return w.header
+}
+
+func (w *writeCountingResponseWriter) Write(p []byte) (int, error) {
+	w.writes++
+	w.body = append(w.body, p...)
+	return len(p), nil
+}
+
+func (w *writeCountingResponseWriter) WriteHeader(statusCode int) {}
+
+func (w *writeCountingResponseWriter) Flush() {}
+
+func TestStreamComponentWritesIncrementally(t *testing.T) {
+	w := &writeCountingResponseWriter{}
+	component := templ.ComponentFunc(func(ctx context.Context, writer io.Writer) error {
+		if _, err := writer.Write([]byte("a")); err != nil {
+			return err
+		}
+		if _, err := writer.Write([]byte("b")); err != nil {
+			return err
+		}
+		return nil
+	})
+
+	if err := streamComponent(context.Background(), w, component); err != nil {
+		t.Fatalf("streamComponent() error = %v", err)
+	}
+	if got := string(w.body); got != "ab" {
+		t.Fatalf("body = %q, want %q", got, "ab")
+	}
+	if w.writes < 2 {
+		t.Fatalf("writes = %d, want at least 2 incremental writes", w.writes)
+	}
+}