Dragonfly2

1
/*
2
 *     Copyright 2022 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
package peer
18

19
import (
20
	"sync"
21
	"testing"
22

23
	testifyassert "github.com/stretchr/testify/assert"
24
)
25

26
func TestBroker(t *testing.T) {
27
	var testCases = []struct {
28
		name   string
29
		total  int32
30
		pubSeq []int32
31
		desire []*PieceInfo
32
	}{
33
		{
34
			name:   "order publish",
35
			total:  3,
36
			pubSeq: []int32{0, 1, 2},
37
			desire: []*PieceInfo{
38
				{
39
					Num:        0,
40
					OrderedNum: 0,
41
					Finished:   false,
42
				},
43
				{
44
					Num:        1,
45
					OrderedNum: 1,
46
					Finished:   false,
47
				},
48
				{
49
					Num:        2,
50
					OrderedNum: 2,
51
					Finished:   true,
52
				},
53
			},
54
		},
55
		{
56
			name:   "partial order publish",
57
			total:  3,
58
			pubSeq: []int32{0, 2, 1},
59
			desire: []*PieceInfo{
60
				{
61
					Num:        0,
62
					OrderedNum: 0,
63
					Finished:   false,
64
				},
65
				{
66
					Num:        2,
67
					OrderedNum: 0,
68
					Finished:   false,
69
				},
70
				{
71
					Num:        1,
72
					OrderedNum: 2,
73
					Finished:   true,
74
				},
75
			},
76
		},
77
		{
78
			name:   "inverted order publish",
79
			total:  3,
80
			pubSeq: []int32{2, 1, 0},
81
			desire: []*PieceInfo{
82
				{
83
					Num:        2,
84
					OrderedNum: -1,
85
					Finished:   false,
86
				},
87
				{
88
					Num:        1,
89
					OrderedNum: -1,
90
					Finished:   false,
91
				},
92
				{
93
					Num:        0,
94
					OrderedNum: 2,
95
					Finished:   true,
96
				},
97
			},
98
		},
99
	}
100
	for _, tc := range testCases {
101
		t.Run(tc.name, func(t *testing.T) {
102
			assert := testifyassert.New(t)
103
			broker := newPieceBroker()
104
			go broker.Start()
105
			ch := broker.Subscribe()
106
			done := make(chan struct{})
107
			var received []*PieceInfo
108

109
			wg := sync.WaitGroup{}
110
			wg.Add(len(tc.pubSeq))
111
			go func() {
112
				for {
113
					select {
114
					case <-done:
115
						return
116
					case info := <-ch:
117
						received = append(received, info)
118
						wg.Done()
119
					}
120
				}
121
			}()
122
			var sent int32
123
			for _, n := range tc.pubSeq {
124
				sent++
125
				var finished bool
126
				if sent == tc.total {
127
					finished = true
128
				}
129
				broker.Publish(&PieceInfo{
130
					Num:      n,
131
					Finished: finished,
132
				})
133
			}
134
			wg.Wait()
135
			broker.Stop()
136
			close(done)
137
			assert.Equal(tc.desire, received)
138
		})
139
	}
140
}
141