package laniakea import ( "context" "errors" "testing" "time" ) type cancelAwareObserver struct { testObserver started chan struct{} } func (o *cancelAwareObserver) OnUpdateReceived(ctx context.Context, _ UpdateReceivedEvent) { close(o.started) <-ctx.Done() } type stubbornObserver struct { testObserver started chan struct{} release chan struct{} } func (o *stubbornObserver) OnUpdateReceived(context.Context, UpdateReceivedEvent) { close(o.started) <-o.release } func TestObserverDispatcherCancelsCallbackDuringClose(t *testing.T) { observer := &cancelAwareObserver{started: make(chan struct{})} dispatcher := newObserverDispatcher(observer, nil) dispatcher.enqueue(context.Background(), UpdateReceivedEvent{}) <-observer.started ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() if err := dispatcher.close(ctx); err != nil { t.Fatalf("close returned error: %v", err) } } func TestObserverDispatcherCloseTimeout(t *testing.T) { observer := &stubbornObserver{started: make(chan struct{}), release: make(chan struct{})} dispatcher := newObserverDispatcher(observer, nil) dispatcher.enqueue(context.Background(), UpdateReceivedEvent{}) <-observer.started ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) defer cancel() if err := dispatcher.close(ctx); !errors.Is(err, ErrObserverShutdownTimeout) { t.Fatalf("close error = %v, want ErrObserverShutdownTimeout", err) } close(observer.release) if err := dispatcher.close(context.Background()); err != nil { t.Fatalf("second close returned error: %v", err) } }