sync from project

middleware/cleanup/cleanup_test.go

@@ -0,0 +1,54 @@
+package cleanup
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+)
+
+func TestCleanupMiddleware(t *testing.T) {
+	var executionOrder []string
+
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		Register(r.Context(), CleanupFunc(func() {
+			executionOrder = append(executionOrder, "first")
+		}))
+		Register(r.Context(), CleanupFunc(func() {
+			executionOrder = append(executionOrder, "second")
+		}))
+		w.WriteHeader(http.StatusOK)
+	})
+
+	req := httptest.NewRequest("GET", "/", nil)
+	rr := httptest.NewRecorder()
+
+	middleware := Collect().WrapHandler(handler)
+	middleware.ServeHTTP(rr, req)
+
+	if rr.Code != http.StatusOK {
+		t.Errorf("handler returned wrong status code: got %v want %v",
+			rr.Code, http.StatusOK)
+	}
+
+	expectedOrder := "second,first"
+	actualOrder := strings.Join(executionOrder, ",")
+	if actualOrder != expectedOrder {
+		t.Errorf("cleanup functions executed in wrong order: got %s want %s",
+			actualOrder, expectedOrder)
+	}
+}
+
+func TestCleanupMissingContext(t *testing.T) {
+	// This test ensures that Register does not panic when the context is missing.
+	// The function should fail silently.
+	defer func() {
+		if r := recover(); r != nil {
+			t.Errorf("The code panicked when it should not have")
+		}
+	}()
+
+	req := httptest.NewRequest("GET", "/", nil)
+	// No middleware, so no context
+	Register(req.Context(), CleanupFunc(func() {}))
+}