131 lines
2.8 KiB
Go
131 lines
2.8 KiB
Go
package messaging
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.pyer.club/kingecg/gotidb/pkg/model"
|
|
"github.com/nats-io/nats.go/jetstream"
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
// 模拟NATS连接
|
|
type mockNATSConn struct {
|
|
closeFunc func() error
|
|
}
|
|
|
|
func (m *mockNATSConn) Close() error {
|
|
if m.closeFunc != nil {
|
|
return m.closeFunc()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// 模拟JetStream
|
|
type mockJetStream struct {
|
|
publishFunc func(ctx context.Context, subject string, data []byte) (jetstream.PubAck, error)
|
|
}
|
|
|
|
func (m *mockJetStream) Publish(ctx context.Context, subject string, data []byte) (jetstream.PubAck, error) {
|
|
if m.publishFunc != nil {
|
|
return m.publishFunc(ctx, subject, data)
|
|
}
|
|
return jetstream.PubAck{}, nil
|
|
}
|
|
|
|
// 模拟Stream
|
|
type mockStream struct {
|
|
createOrUpdateConsumerFunc func(ctx context.Context, cfg jetstream.ConsumerConfig) (jetstream.Consumer, error)
|
|
}
|
|
|
|
func (m *mockStream) CreateOrUpdateConsumer(ctx context.Context, cfg jetstream.ConsumerConfig) (jetstream.Consumer, error) {
|
|
if m.createOrUpdateConsumerFunc != nil {
|
|
return m.createOrUpdateConsumerFunc(ctx, cfg)
|
|
}
|
|
return nil, nil
|
|
}
|
|
|
|
// 模拟Consumer
|
|
type mockConsumer struct {
|
|
messagesFunc func() (jetstream.MessagesContext, error)
|
|
}
|
|
|
|
func (m *mockConsumer) Messages() (jetstream.MessagesContext, error) {
|
|
if m.messagesFunc != nil {
|
|
return m.messagesFunc()
|
|
}
|
|
return nil, nil
|
|
}
|
|
|
|
func TestNATSMessaging_Publish(t *testing.T) {
|
|
publishCalled := false
|
|
mockJS := &mockJetStream{
|
|
publishFunc: func(ctx context.Context, subject string, data []byte) (jetstream.PubAck, error) {
|
|
publishCalled = true
|
|
return jetstream.PubAck{}, nil
|
|
},
|
|
}
|
|
|
|
messaging := &NATSMessaging{
|
|
conn: &mockNATSConn{},
|
|
js: mockJS,
|
|
}
|
|
|
|
id := model.DataPointID{
|
|
DeviceID: "device1",
|
|
MetricCode: "metric1",
|
|
Labels: map[string]string{"env": "test"},
|
|
}
|
|
value := model.DataValue{
|
|
Timestamp: time.Now(),
|
|
Value: 42.0,
|
|
}
|
|
|
|
err := messaging.Publish(context.Background(), id, value)
|
|
assert.NoError(t, err)
|
|
assert.True(t, publishCalled)
|
|
}
|
|
|
|
func TestNATSMessaging_Subscribe(t *testing.T) {
|
|
handlerCalled := false
|
|
handler := func(msg DataMessage) error {
|
|
handlerCalled = true
|
|
return nil
|
|
}
|
|
|
|
mockConsumer := &mockConsumer{}
|
|
mockStream := &mockStream{
|
|
createOrUpdateConsumerFunc: func(ctx context.Context, cfg jetstream.ConsumerConfig) (jetstream.Consumer, error) {
|
|
return mockConsumer, nil
|
|
},
|
|
}
|
|
|
|
messaging := &NATSMessaging{
|
|
conn: &mockNATSConn{},
|
|
stream: mockStream,
|
|
}
|
|
|
|
err := messaging.Subscribe(handler)
|
|
assert.NoError(t, err)
|
|
assert.Contains(t, messaging.handlers, handler)
|
|
}
|
|
|
|
func TestNATSMessaging_Close(t *testing.T) {
|
|
closeCalled := false
|
|
mockConn := &mockNATSConn{
|
|
closeFunc: func() error {
|
|
closeCalled = true
|
|
return nil
|
|
},
|
|
}
|
|
|
|
messaging := &NATSMessaging{
|
|
conn: mockConn,
|
|
}
|
|
|
|
err := messaging.Close()
|
|
assert.NoError(t, err)
|
|
assert.True(t, closeCalled)
|
|
}
|