@@ -48,19 +48,21 @@ func newTestComparator() *frameworkmocks.MockItemComparator {
4848 }
4949}
5050
51- func newTestBand (queues map [ string ] framework.FlowQueueAccessor ) * frameworkmocks.MockPriorityBandAccessor {
51+ func newTestBand (queues ... framework.FlowQueueAccessor ) * frameworkmocks.MockPriorityBandAccessor {
5252 flowIDs := make ([]string , 0 , len (queues ))
53- for id := range queues {
54- flowIDs = append (flowIDs , id )
53+ queuesByID := make (map [string ]framework.FlowQueueAccessor , len (queues ))
54+ for _ , q := range queues {
55+ flowIDs = append (flowIDs , q .FlowSpec ().ID )
56+ queuesByID [q .FlowSpec ().ID ] = q
5557 }
5658 return & frameworkmocks.MockPriorityBandAccessor {
5759 FlowIDsFunc : func () []string { return flowIDs },
5860 QueueFunc : func (id string ) framework.FlowQueueAccessor {
59- return queues [id ]
61+ return queuesByID [id ]
6062 },
6163 IterateQueuesFunc : func (iterator func (queue framework.FlowQueueAccessor ) bool ) {
6264 for _ , id := range flowIDs {
63- if ! iterator (queues [id ]) {
65+ if ! iterator (queuesByID [id ]) {
6466 break
6567 }
6668 }
@@ -111,90 +113,88 @@ func TestBestHead_SelectQueue(t *testing.T) {
111113 shouldPanic bool
112114 }{
113115 {
114- name : "BasicSelection_TwoQueues" ,
115- band : newTestBand (map [string ]framework.FlowQueueAccessor {
116- flow1 : queue1 ,
117- flow2 : queue2 ,
118- }),
116+ name : "BasicSelection_TwoQueues" ,
117+ band : newTestBand (queue1 , queue2 ),
119118 expectedQueueID : flow1 ,
120119 },
121120 {
122- name : "IgnoresEmptyQueues" ,
123- band : newTestBand (map [string ]framework.FlowQueueAccessor {
124- flow1 : queue1 ,
125- "flowEmpty" : queueEmpty ,
126- flow2 : queue2 ,
127- }),
121+ name : "IgnoresEmptyQueues" ,
122+ band : newTestBand (queue1 , queueEmpty , queue2 ),
128123 expectedQueueID : flow1 ,
129124 },
130125 {
131126 name : "SingleNonEmptyQueue" ,
132- band : newTestBand (map [ string ]framework. FlowQueueAccessor { flow1 : queue1 } ),
127+ band : newTestBand (queue1 ),
133128 expectedQueueID : flow1 ,
134129 },
135130 {
136131 name : "ComparatorCompatibility" ,
137- band : newTestBand (map [ string ]framework. FlowQueueAccessor {
138- flow1 : & frameworkmocks.MockFlowQueueAccessor {
132+ band : newTestBand (
133+ & frameworkmocks.MockFlowQueueAccessor {
139134 LenV : 1 ,
140135 PeekHeadV : itemBetter ,
141136 FlowSpecV : types.FlowSpecification {ID : flow1 },
142137 ComparatorV : & frameworkmocks.MockItemComparator {ScoreTypeV : "typeA" , FuncV : enqueueTimeComparatorFunc },
143138 },
144- flow2 : & frameworkmocks.MockFlowQueueAccessor {
139+ & frameworkmocks.MockFlowQueueAccessor {
145140 LenV : 1 ,
146141 PeekHeadV : itemWorse ,
147142 FlowSpecV : types.FlowSpecification {ID : flow2 },
148143 ComparatorV : & frameworkmocks.MockItemComparator {ScoreTypeV : "typeB" , FuncV : enqueueTimeComparatorFunc },
149144 },
150- } ),
145+ ),
151146 expectedErr : framework .ErrIncompatiblePriorityType ,
152147 },
153148 {
154149 name : "QueuePeekHeadErrors" ,
155- band : newTestBand (map [ string ]framework. FlowQueueAccessor {
156- flow1 : & frameworkmocks.MockFlowQueueAccessor {
150+ band : newTestBand (
151+ & frameworkmocks.MockFlowQueueAccessor {
157152 LenV : 1 ,
158153 PeekHeadErrV : errors .New ("peek error" ),
159154 FlowSpecV : types.FlowSpecification {ID : flow1 },
160155 ComparatorV : newTestComparator (),
161156 },
162- flow2 : queue2 ,
163- } ),
157+ queue2 ,
158+ ),
164159 expectedQueueID : flow2 ,
165160 },
166161 {
167162 name : "QueueComparatorIsNil" ,
168- band : newTestBand (map [ string ]framework. FlowQueueAccessor {
169- flow1 : & frameworkmocks.MockFlowQueueAccessor {
163+ band : newTestBand (
164+ & frameworkmocks.MockFlowQueueAccessor {
170165 LenV : 1 ,
171166 PeekHeadV : itemBetter ,
172167 FlowSpecV : types.FlowSpecification {ID : flow1 },
173168 ComparatorV : nil ,
174169 },
175- flow2 : queue2 ,
176- } ),
170+ queue2 ,
171+ ),
177172 shouldPanic : true ,
178173 },
179174 {
180175 name : "ComparatorFuncIsNil" ,
181- band : newTestBand (map [ string ]framework. FlowQueueAccessor {
182- flow1 : & frameworkmocks.MockFlowQueueAccessor {
176+ band : newTestBand (
177+ & frameworkmocks.MockFlowQueueAccessor {
183178 LenV : 1 ,
184179 PeekHeadV : itemBetter ,
185180 FlowSpecV : types.FlowSpecification {ID : flow1 },
186181 ComparatorV : & frameworkmocks.MockItemComparator {ScoreTypeV : commonScoreType , FuncV : nil },
187182 },
188- flow2 : queue2 ,
189- } ),
183+ queue2 ,
184+ ),
190185 shouldPanic : true ,
191186 },
192187 {
193188 name : "AllQueuesEmpty" ,
194- band : newTestBand (map [string ]framework.FlowQueueAccessor {
195- "empty1" : queueEmpty ,
196- "empty2" : queueEmpty ,
197- }),
189+ band : newTestBand (
190+ queueEmpty ,
191+ & frameworkmocks.MockFlowQueueAccessor {
192+ LenV : 0 ,
193+ PeekHeadErrV : framework .ErrQueueEmpty ,
194+ FlowSpecV : types.FlowSpecification {ID : "flowEmpty2" },
195+ ComparatorV : newTestComparator (),
196+ },
197+ ),
198198 },
199199 {
200200 name : "NilBand" ,
0 commit comments