github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/mq/mq_test.go (about) 1 package mq 2 3 import ( 4 "github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc" 5 "github.com/stretchr/testify/assert" 6 "sync" 7 "testing" 8 "time" 9 ) 10 11 func TestPublishAndSubscribe(t *testing.T) { 12 t.Parallel() 13 asserts := assert.New(t) 14 mq := NewMQ() 15 16 // No subscriber 17 { 18 asserts.NotPanics(func() { 19 mq.Publish("No subscriber", Message{}) 20 }) 21 } 22 23 // One channel subscriber 24 { 25 topic := "One channel subscriber" 26 msg := Message{TriggeredBy: "Tester"} 27 notifier := mq.Subscribe(topic, 0) 28 mq.Publish(topic, msg) 29 wg := sync.WaitGroup{} 30 wg.Add(1) 31 go func() { 32 wg.Done() 33 msgRecv := <-notifier 34 asserts.Equal(msg, msgRecv) 35 }() 36 wg.Wait() 37 } 38 39 // two channel subscriber 40 { 41 topic := "two channel subscriber" 42 msg := Message{TriggeredBy: "Tester"} 43 notifier := mq.Subscribe(topic, 0) 44 notifier2 := mq.Subscribe(topic, 0) 45 mq.Publish(topic, msg) 46 wg := sync.WaitGroup{} 47 wg.Add(2) 48 go func() { 49 wg.Done() 50 msgRecv := <-notifier 51 asserts.Equal(msg, msgRecv) 52 }() 53 go func() { 54 wg.Done() 55 msgRecv := <-notifier2 56 asserts.Equal(msg, msgRecv) 57 }() 58 wg.Wait() 59 } 60 61 // two channel subscriber, one timeout 62 { 63 topic := "two channel subscriber, one timeout" 64 msg := Message{TriggeredBy: "Tester"} 65 mq.Subscribe(topic, 0) 66 notifier2 := mq.Subscribe(topic, 0) 67 mq.Publish(topic, msg) 68 wg := sync.WaitGroup{} 69 wg.Add(1) 70 go func() { 71 wg.Done() 72 msgRecv := <-notifier2 73 asserts.Equal(msg, msgRecv) 74 }() 75 wg.Wait() 76 } 77 78 // two channel subscriber, one unsubscribe 79 { 80 topic := "two channel subscriber, one unsubscribe" 81 msg := Message{TriggeredBy: "Tester"} 82 mq.Subscribe(topic, 0) 83 notifier2 := mq.Subscribe(topic, 0) 84 notifier := mq.Subscribe(topic, 0) 85 mq.Unsubscribe(topic, notifier) 86 mq.Publish(topic, msg) 87 wg := sync.WaitGroup{} 88 wg.Add(1) 89 go func() { 90 wg.Done() 91 msgRecv := <-notifier2 92 asserts.Equal(msg, msgRecv) 93 }() 94 wg.Wait() 95 96 select { 97 case <-notifier: 98 t.Error() 99 default: 100 } 101 } 102 } 103 104 func TestAria2Interface(t *testing.T) { 105 t.Parallel() 106 asserts := assert.New(t) 107 mq := NewMQ() 108 var ( 109 OnDownloadStart int 110 OnDownloadPause int 111 OnDownloadStop int 112 OnDownloadComplete int 113 OnDownloadError int 114 ) 115 l := sync.Mutex{} 116 117 mq.SubscribeCallback("TestAria2Interface", func(message Message) { 118 asserts.Equal("TestAria2Interface", message.TriggeredBy) 119 l.Lock() 120 defer l.Unlock() 121 switch message.Event { 122 case "1": 123 OnDownloadStart++ 124 case "2": 125 OnDownloadPause++ 126 case "5": 127 OnDownloadStop++ 128 case "4": 129 OnDownloadComplete++ 130 case "3": 131 OnDownloadError++ 132 } 133 }) 134 135 mq.OnDownloadStart([]rpc.Event{{"TestAria2Interface"}, {"TestAria2Interface"}}) 136 mq.OnDownloadPause([]rpc.Event{{"TestAria2Interface"}, {"TestAria2Interface"}}) 137 mq.OnDownloadStop([]rpc.Event{{"TestAria2Interface"}, {"TestAria2Interface"}}) 138 mq.OnDownloadComplete([]rpc.Event{{"TestAria2Interface"}, {"TestAria2Interface"}}) 139 mq.OnDownloadError([]rpc.Event{{"TestAria2Interface"}, {"TestAria2Interface"}}) 140 mq.OnBtDownloadComplete([]rpc.Event{{"TestAria2Interface"}, {"TestAria2Interface"}}) 141 142 time.Sleep(time.Duration(500) * time.Millisecond) 143 144 asserts.Equal(2, OnDownloadStart) 145 asserts.Equal(2, OnDownloadPause) 146 asserts.Equal(2, OnDownloadStop) 147 asserts.Equal(4, OnDownloadComplete) 148 asserts.Equal(2, OnDownloadError) 149 }