@@ -159,6 +159,40 @@ func TestAllocate(t *testing.T) {
159
159
require .Equal (t , resp , respOut )
160
160
}
161
161
162
+ func TestGetPreferredAllocation (t * testing.T ) {
163
+ socket := path .Join ("/tmp" , esocketName )
164
+ callbackCount := 0
165
+ callbackChan := make (chan int )
166
+ p , e := esetup (t , []* pluginapi.Device {}, socket , "mock" , func (n string , d []pluginapi.Device ) {
167
+ callbackCount ++
168
+ callbackChan <- callbackCount
169
+ })
170
+ defer ecleanup (t , p , e )
171
+
172
+ resp := & pluginapi.PreferredAllocationResponse {
173
+ ContainerResponses : []* pluginapi.ContainerPreferredAllocationResponse {
174
+ {DeviceIDs : []string {"device0" , "device1" , "device2" }},
175
+ },
176
+ }
177
+
178
+ p .SetGetPreferredAllocFunc (func (r * pluginapi.PreferredAllocationRequest , devs map [string ]pluginapi.Device ) (* pluginapi.PreferredAllocationResponse , error ) {
179
+ return resp , nil
180
+ })
181
+
182
+ go e .run ()
183
+ // Wait for the callback to be issued.
184
+ select {
185
+ case <- callbackChan :
186
+ break
187
+ case <- time .After (time .Second ):
188
+ t .FailNow ()
189
+ }
190
+
191
+ respOut , err := e .getPreferredAllocation ([]string {}, []string {}, - 1 )
192
+ require .NoError (t , err )
193
+ require .Equal (t , resp , respOut )
194
+ }
195
+
162
196
func esetup (t * testing.T , devs []* pluginapi.Device , socket , resourceName string , callback monitorCallback ) (* Stub , * endpointImpl ) {
163
197
p := NewDevicePluginStub (devs , socket , resourceName , false , false )
164
198
0 commit comments