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  }