milvus/pkg/streaming/util/message/message_handler_test.go
Zhen Ye f0f5147aef
fix: streaming consumer may get stucked when handler is un-consumed (#36818)
issue: #36378

Signed-off-by: chyezh <chyezh@outlook.com>
2024-10-14 15:23:23 +08:00

28 lines
513 B
Go

package message
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
)
func TestMessageHandler(t *testing.T) {
ch := make(chan ImmutableMessage, 1)
h := ChanMessageHandler(ch)
ok, err := h.Handle(context.Background(), nil)
assert.NoError(t, err)
assert.True(t, ok)
ctx, cancel := context.WithCancel(context.Background())
cancel()
ok, err = h.Handle(ctx, nil)
assert.ErrorIs(t, err, ctx.Err())
assert.False(t, ok)
assert.Nil(t, <-ch)
h.Close()
_, ok = <-ch
assert.False(t, ok)
}