...
1 package worker
2
3 import (
4 "os"
5 "testing"
6 "time"
7
8 "github.com/google/uuid"
9 "github.com/stretchr/testify/mock"
10 "github.com/stretchr/testify/require"
11
12 "edge-infra.dev/pkg/edge/datasync/internal/config"
13
14 "edge-infra.dev/pkg/edge/datasync/chirp/model"
15 "edge-infra.dev/pkg/edge/datasync/internal/config/fake"
16 )
17
18 type MessageProviderMock struct{ mock.Mock }
19 type MessageSenderMock struct{ mock.Mock }
20
21 func TestDoWork(t *testing.T) {
22 fake.SetEnv()
23
24 msgProviderMock := new(MessageProviderMock)
25
26 messages := make(map[string]model.Message)
27 msgID1 := uuid.New()
28 messages[msgID1.String()] = model.Message{ID: msgID1, Type: "tlog", FilePath: "testdata/xxx"}
29 msgID2 := uuid.New()
30 messages[msgID2.String()] = model.Message{ID: msgID2, Type: "tlog", FilePath: "testdata/xxx"}
31
32 msgProviderMock.On("GetMessages").Return(messages)
33
34 sentMessages := []model.Message{
35 messages[msgID1.String()],
36 messages[msgID2.String()],
37 }
38 msgSenderMock := new(MessageSenderMock)
39 msgSenderMock.On("Send", messages).Return(sentMessages)
40 cfg := &config.Config{
41 RunMessageWorker: true,
42 WorkerPolling: 6,
43 }
44 worker := NewMessageWorker(1, msgProviderMock, msgSenderMock, cfg)
45
46 require.NotEmpty(t, worker)
47
48 go worker.DoWork()
49
50 time.Sleep(100 * time.Millisecond)
51 os.Setenv("RUN_MESSAGE_WORKER", "false")
52
53 msgProviderMock.AssertExpectations(t)
54 msgSenderMock.AssertExpectations(t)
55 }
56
57 func (p *MessageProviderMock) GetMessages() map[string]model.Message {
58 args := p.Called()
59
60 result := args.Get(0)
61 return result.(map[string]model.Message)
62 }
63
64 func (s *MessageSenderMock) Send(messages map[string]model.Message) []model.Message {
65 args := s.Called(messages)
66 result := args.Get(0)
67
68 return result.([]model.Message)
69 }
70
View as plain text