mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-09 23:40:06 +08:00
333a0c8a53
Signed-off-by: xiaofan-luan <xiaofan.luan@zilliz.com>
253 lines
6.9 KiB
Go
253 lines
6.9 KiB
Go
// Licensed to the LF AI & Data foundation under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package rmq
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math/rand"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/apache/pulsar-client-go/pulsar"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
rocksmqimplclient "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/client"
|
|
rocksmqimplserver "github.com/milvus-io/milvus/internal/mq/mqimpl/rocksmq/server"
|
|
"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
|
|
pulsarwrapper "github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper/pulsar"
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
|
)
|
|
|
|
func TestMain(m *testing.M) {
|
|
paramtable.Init()
|
|
pt := paramtable.Get()
|
|
pt.Save(pt.ServiceParam.MQCfg.EnablePursuitMode.Key, "false")
|
|
|
|
rand.Seed(time.Now().UnixNano())
|
|
path := "/tmp/milvus/rdb_data"
|
|
defer os.RemoveAll(path)
|
|
paramtable.Get().Save("rocksmq.compressionTypes", "0,0,0,0,0")
|
|
_ = rocksmqimplserver.InitRocksMQ(path)
|
|
exitCode := m.Run()
|
|
defer rocksmqimplserver.CloseRocksMQ()
|
|
os.Exit(exitCode)
|
|
}
|
|
|
|
func Test_NewRmqClient(t *testing.T) {
|
|
client, err := createRmqClient()
|
|
defer client.Close()
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, client)
|
|
}
|
|
|
|
func TestRmqClient_CreateProducer(t *testing.T) {
|
|
opts := rocksmqimplclient.Options{}
|
|
client, err := NewClient(opts)
|
|
defer client.Close()
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, client)
|
|
|
|
topic := "TestRmqClient_CreateProducer"
|
|
proOpts := mqwrapper.ProducerOptions{Topic: topic}
|
|
producer, err := client.CreateProducer(proOpts)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, producer)
|
|
|
|
defer producer.Close()
|
|
|
|
rmqProducer := producer.(*rmqProducer)
|
|
defer rmqProducer.Close()
|
|
assert.Equal(t, rmqProducer.Topic(), topic)
|
|
|
|
msg := &mqwrapper.ProducerMessage{
|
|
Payload: []byte{},
|
|
Properties: nil,
|
|
}
|
|
_, err = rmqProducer.Send(context.TODO(), msg)
|
|
assert.NoError(t, err)
|
|
|
|
invalidOpts := mqwrapper.ProducerOptions{Topic: ""}
|
|
producer, e := client.CreateProducer(invalidOpts)
|
|
assert.Nil(t, producer)
|
|
assert.Error(t, e)
|
|
}
|
|
|
|
func TestRmqClient_GetLatestMsg(t *testing.T) {
|
|
client, err := createRmqClient()
|
|
assert.NoError(t, err)
|
|
defer client.Close()
|
|
|
|
topic := fmt.Sprintf("t2GetLatestMsg-%d", rand.Int())
|
|
proOpts := mqwrapper.ProducerOptions{Topic: topic}
|
|
producer, err := client.CreateProducer(proOpts)
|
|
assert.NoError(t, err)
|
|
defer producer.Close()
|
|
|
|
for i := 0; i < 10; i++ {
|
|
msg := &mqwrapper.ProducerMessage{
|
|
Payload: []byte{byte(i)},
|
|
Properties: nil,
|
|
}
|
|
_, err = producer.Send(context.TODO(), msg)
|
|
assert.NoError(t, err)
|
|
}
|
|
|
|
subName := "subName"
|
|
consumerOpts := mqwrapper.ConsumerOptions{
|
|
Topic: topic,
|
|
SubscriptionName: subName,
|
|
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
|
BufSize: 1024,
|
|
}
|
|
|
|
consumer, err := client.Subscribe(consumerOpts)
|
|
assert.NoError(t, err)
|
|
|
|
expectLastMsg, err := consumer.GetLatestMsgID()
|
|
assert.NoError(t, err)
|
|
|
|
var actualLastMsg mqwrapper.Message
|
|
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
|
defer cancel()
|
|
for i := 0; i < 10; i++ {
|
|
select {
|
|
case <-ctx.Done():
|
|
fmt.Println(i)
|
|
assert.FailNow(t, "consumer failed to yield message in 100 milliseconds")
|
|
case msg := <-consumer.Chan():
|
|
consumer.Ack(msg)
|
|
actualLastMsg = msg
|
|
}
|
|
}
|
|
require.NotNil(t, actualLastMsg)
|
|
ret, err := expectLastMsg.LessOrEqualThan(actualLastMsg.ID().Serialize())
|
|
assert.NoError(t, err)
|
|
assert.True(t, ret)
|
|
}
|
|
|
|
func TestRmqClient_Subscribe(t *testing.T) {
|
|
client, err := createRmqClient()
|
|
defer client.Close()
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, client)
|
|
|
|
topic := "TestRmqClient_Subscribe"
|
|
proOpts := mqwrapper.ProducerOptions{Topic: topic}
|
|
producer, err := client.CreateProducer(proOpts)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, producer)
|
|
defer producer.Close()
|
|
|
|
subName := "subName"
|
|
consumerOpts := mqwrapper.ConsumerOptions{
|
|
Topic: subName,
|
|
SubscriptionName: subName,
|
|
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
|
BufSize: 0,
|
|
}
|
|
consumer, err := client.Subscribe(consumerOpts)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, consumer)
|
|
|
|
consumerOpts = mqwrapper.ConsumerOptions{
|
|
Topic: "",
|
|
SubscriptionName: subName,
|
|
SubscriptionInitialPosition: mqwrapper.SubscriptionPositionEarliest,
|
|
BufSize: 1024,
|
|
}
|
|
|
|
consumer, err = client.Subscribe(consumerOpts)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, consumer)
|
|
|
|
consumerOpts.Topic = topic
|
|
consumer, err = client.Subscribe(consumerOpts)
|
|
defer consumer.Close()
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, consumer)
|
|
assert.Equal(t, consumer.Subscription(), subName)
|
|
|
|
msg := &mqwrapper.ProducerMessage{
|
|
Payload: []byte{1},
|
|
Properties: nil,
|
|
}
|
|
_, err = producer.Send(context.TODO(), msg)
|
|
assert.NoError(t, err)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
|
defer cancel()
|
|
select {
|
|
case <-ctx.Done():
|
|
assert.FailNow(t, "consumer failed to yield message in 100 milliseconds")
|
|
case msg := <-consumer.Chan():
|
|
consumer.Ack(msg)
|
|
rmqmsg := msg.(*rocksmqimplclient.RmqMessage)
|
|
msgPayload := rmqmsg.Payload()
|
|
assert.NotEmpty(t, msgPayload)
|
|
msgTopic := rmqmsg.Topic()
|
|
assert.Equal(t, msgTopic, topic)
|
|
msgProp := rmqmsg.Properties()
|
|
assert.Empty(t, msgProp)
|
|
msgID := rmqmsg.ID()
|
|
rID := msgID.(*rocksmqimplserver.RmqID)
|
|
assert.NotZero(t, rID)
|
|
}
|
|
}
|
|
|
|
func TestRmqClient_EarliestMessageID(t *testing.T) {
|
|
client, _ := createRmqClient()
|
|
defer client.Close()
|
|
|
|
mid := client.EarliestMessageID()
|
|
assert.NotNil(t, mid)
|
|
}
|
|
|
|
func TestRmqClient_StringToMsgID(t *testing.T) {
|
|
client, _ := createRmqClient()
|
|
defer client.Close()
|
|
|
|
str := "5"
|
|
res, err := client.StringToMsgID(str)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, res)
|
|
|
|
str = "X"
|
|
res, err = client.StringToMsgID(str)
|
|
assert.Nil(t, res)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestRmqClient_BytesToMsgID(t *testing.T) {
|
|
client, _ := createRmqClient()
|
|
defer client.Close()
|
|
|
|
mid := pulsar.EarliestMessageID()
|
|
binary := pulsarwrapper.SerializePulsarMsgID(mid)
|
|
|
|
res, err := client.BytesToMsgID(binary)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, res)
|
|
}
|
|
|
|
func createRmqClient() (*rmqClient, error) {
|
|
opts := rocksmqimplclient.Options{}
|
|
return NewClient(opts)
|
|
}
|