manager

func Register[T any](r *Relay, topic string, fn Handler[T])

{
	wrapper := func(raw []byte) error {
		r.ctxMu.RLock()
		ctx := r.ctx
		r.ctxMu.RUnlock()
		if ctx == nil {
			ctx = context.Background()
		}

		var payload T
		if err := json.Unmarshal(raw, &payload); err != nil {
			return fmt.Errorf("payload unmarshal failed: %w", err)
		}
		return fn(ctx, payload)
	}

	r.handlerMu.Lock()
	r.handlers[topic] = wrapper
	r.handlerMu.Unlock()
}

func TestHandlerReceivesStartContext(t *testing.T)

{
	r := New()

	var receivedCtx context.Context
	var mu sync.Mutex

	Register(r, "test-topic", func(ctx context.Context, payload string) error {
		mu.Lock()
		receivedCtx = ctx
		mu.Unlock()
		return nil
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	done := make(chan error, 1)
	go func() {
		done <- r.Start(ctx)
	}()

	// Publish a message and wait for handler to run.
	if err := Enqueue(ctx, r, "test-topic", "hello"); err != nil {
		t.Fatalf("Enqueue failed: %v", err)
	}

	// Give the handler time to run.
	time.Sleep(50 * time.Millisecond)

	mu.Lock()
	rc := receivedCtx
	mu.Unlock()

	if rc == nil {
		t.Fatal("handler was not called")
	}
	if rc == context.Background() {
		t.Fatal("handler received context.Background() instead of Start's context")
	}

	// Cancel Start's context and verify it is the same context the handler saw.
	cancel()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("Start did not return after context cancel")
	}

	// The handler's context should also be done now.
	select {
	case <-rc.Done():
	default:
		t.Fatal("handler context is not done after Start's context was cancelled")
	}
}

func TestRegisterAfterStartDoesNotDeadlock(t *testing.T)

{
	r := New()

	Register(r, "initial-topic", func(ctx context.Context, payload string) error {
		return nil
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	startDone := make(chan error, 1)
	go func() {
		startDone <- r.Start(ctx)
	}()

	// Give Start time to enter its subscribe loop and block on <-ctx.Done().
	time.Sleep(50 * time.Millisecond)

	// Register must not deadlock; complete within the timeout.
	registered := make(chan struct{})
	go func() {
		Register(r, "late-topic", func(ctx context.Context, payload string) error {
			return nil
		})
		close(registered)
	}()

	select {
	case <-registered:
		// success
	case <-time.After(time.Second):
		t.Fatal("Register called after Start deadlocked")
	}

	cancel()
	select {
	case <-startDone:
	case <-time.After(time.Second):
		t.Fatal("Start did not return after context cancel")
	}
}