first commit

transports/bifrost-http/handlers/wsresponses_test.go

@@ -0,0 +1,68 @@
+package handlers
+
+import (
+	"testing"
+
+	"github.com/maximhq/bifrost/core/schemas"
+	"github.com/maximhq/bifrost/framework/kvstore"
+	"github.com/maximhq/bifrost/framework/logstore"
+	"github.com/maximhq/bifrost/transports/bifrost-http/lib"
+)
+
+type testWSHandlerStore struct {
+	allowDirectKeys bool
+}
+
+func (s testWSHandlerStore) ShouldAllowDirectKeys() bool {
+	return s.allowDirectKeys
+}
+
+func (s testWSHandlerStore) GetHeaderMatcher() *lib.HeaderMatcher {
+	return nil
+}
+
+func (s testWSHandlerStore) GetAvailableProviders() []schemas.ModelProvider {
+	return nil
+}
+
+func (s testWSHandlerStore) GetStreamChunkInterceptor() lib.StreamChunkInterceptor {
+	return nil
+}
+
+func (s testWSHandlerStore) GetAsyncJobExecutor() *logstore.AsyncJobExecutor {
+	return nil
+}
+
+func (s testWSHandlerStore) GetAsyncJobResultTTL() int {
+	return 0
+}
+
+func (s testWSHandlerStore) GetKVStore() *kvstore.Store {
+	return nil
+}
+
+func (s testWSHandlerStore) GetMCPHeaderCombinedAllowlist() schemas.WhiteList {
+	return nil
+}
+
+func TestCreateBifrostContextFromAuth_BaggageSessionIDSetsGrouping(t *testing.T) {
+	ctx, cancel := createBifrostContextFromAuth(testWSHandlerStore{}, &authHeaders{
+		baggage: "foo=bar, session-id=rt-ws-123, baz=qux",
+	})
+	defer cancel()
+
+	if got, _ := ctx.Value(schemas.BifrostContextKeyParentRequestID).(string); got != "rt-ws-123" {
+		t.Fatalf("parent request id = %q, want %q", got, "rt-ws-123")
+	}
+}
+
+func TestCreateBifrostContextFromAuth_EmptyBaggageSessionIDIgnored(t *testing.T) {
+	ctx, cancel := createBifrostContextFromAuth(testWSHandlerStore{}, &authHeaders{
+		baggage: "session-id=   ",
+	})
+	defer cancel()
+
+	if got := ctx.Value(schemas.BifrostContextKeyParentRequestID); got != nil {
+		t.Fatalf("parent request id should be unset, got %#v", got)
+	}
+}