Update existing tests in support of GetPreferredallocation()

pkg/kubelet/cm/devicemanager/endpoint_test.go

@@ -160,7 +160,7 @@ func TestAllocate(t *testing.T) {
 }
 
 func esetup(t *testing.T, devs []*pluginapi.Device, socket, resourceName string, callback monitorCallback) (*Stub, *endpointImpl) {
-	p := NewDevicePluginStub(devs, socket, resourceName, false)
+	p := NewDevicePluginStub(devs, socket, resourceName, false, false)
 
 	err := p.Start()
 	require.NoError(t, err)

pkg/kubelet/cm/devicemanager/manager_test.go

@@ -103,6 +103,7 @@ func TestDevicePluginReRegistration(t *testing.T) {
 		{ID: "Dev3", Health: pluginapi.Healthy},
 	}
 	for _, preStartContainerFlag := range []bool{false, true} {
+		for _, getPreferredAllocationFlag := range []bool{false, true} {
 			m, ch, p1 := setup(t, devs, nil, socketName, pluginSocketName)
 			p1.Register(socketName, testResourceName, "")
 
@@ -117,7 +118,7 @@ func TestDevicePluginReRegistration(t *testing.T) {
 			require.Equal(t, resourceCapacity.Value(), resourceAllocatable.Value(), "capacity should equal to allocatable")
 			require.Equal(t, int64(2), resourceAllocatable.Value(), "Devices are not updated.")
 
-		p2 := NewDevicePluginStub(devs, pluginSocketName+".new", testResourceName, preStartContainerFlag)
+			p2 := NewDevicePluginStub(devs, pluginSocketName+".new", testResourceName, preStartContainerFlag, getPreferredAllocationFlag)
 			err = p2.Start()
 			require.NoError(t, err)
 			p2.Register(socketName, testResourceName, "")
@@ -134,7 +135,7 @@ func TestDevicePluginReRegistration(t *testing.T) {
 			require.Equal(t, int64(2), resourceAllocatable.Value(), "Devices shouldn't change.")
 
 			// Test the scenario that a plugin re-registers with different devices.
-		p3 := NewDevicePluginStub(devsForRegistration, pluginSocketName+".third", testResourceName, preStartContainerFlag)
+			p3 := NewDevicePluginStub(devsForRegistration, pluginSocketName+".third", testResourceName, preStartContainerFlag, getPreferredAllocationFlag)
 			err = p3.Start()
 			require.NoError(t, err)
 			p3.Register(socketName, testResourceName, "")
@@ -153,6 +154,7 @@ func TestDevicePluginReRegistration(t *testing.T) {
 			p3.Stop()
 			cleanup(t, m, p1)
 		}
+	}
 }
 
 // Tests that the device plugin manager correctly handles registration and re-registration by
@@ -186,7 +188,7 @@ func TestDevicePluginReRegistrationProbeMode(t *testing.T) {
 	require.Equal(t, resourceCapacity.Value(), resourceAllocatable.Value(), "capacity should equal to allocatable")
 	require.Equal(t, int64(2), resourceAllocatable.Value(), "Devices are not updated.")
 
-	p2 := NewDevicePluginStub(devs, pluginSocketName+".new", testResourceName, false)
+	p2 := NewDevicePluginStub(devs, pluginSocketName+".new", testResourceName, false, false)
 	err = p2.Start()
 	require.NoError(t, err)
 	// Wait for the second callback to be issued.
@@ -203,7 +205,7 @@ func TestDevicePluginReRegistrationProbeMode(t *testing.T) {
 	require.Equal(t, int64(2), resourceAllocatable.Value(), "Devices are not updated.")
 
 	// Test the scenario that a plugin re-registers with different devices.
-	p3 := NewDevicePluginStub(devsForRegistration, pluginSocketName+".third", testResourceName, false)
+	p3 := NewDevicePluginStub(devsForRegistration, pluginSocketName+".third", testResourceName, false, false)
 	err = p3.Start()
 	require.NoError(t, err)
 	// Wait for the third callback to be issued.
@@ -249,7 +251,7 @@ func setupDeviceManager(t *testing.T, devs []*pluginapi.Device, callback monitor
 }
 
 func setupDevicePlugin(t *testing.T, devs []*pluginapi.Device, pluginSocketName string) *Stub {
-	p := NewDevicePluginStub(devs, pluginSocketName, testResourceName, false)
+	p := NewDevicePluginStub(devs, pluginSocketName, testResourceName, false, false)
 	err := p.Start()
 	require.NoError(t, err)
 	return p
@@ -549,6 +551,7 @@ func (a *activePodsStub) updateActivePods(newPods []*v1.Pod) {
 }
 
 type MockEndpoint struct {
+	getPreferredAllocationFunc func(available, mustInclude []string, size int) (*pluginapi.PreferredAllocationResponse, error)
 	allocateFunc               func(devs []string) (*pluginapi.AllocateResponse, error)
 	initChan                   chan []string
 }
@@ -563,6 +566,13 @@ func (m *MockEndpoint) preStartContainer(devs []string) (*pluginapi.PreStartCont
 	return &pluginapi.PreStartContainerResponse{}, nil
 }
 
+func (m *MockEndpoint) getPreferredAllocation(available, mustInclude []string, size int) (*pluginapi.PreferredAllocationResponse, error) {
+	if m.getPreferredAllocationFunc != nil {
+		return m.getPreferredAllocationFunc(available, mustInclude, size)
+	}
+	return nil, nil
+}
+
 func (m *MockEndpoint) allocate(devs []string) (*pluginapi.AllocateResponse, error) {
 	if m.allocateFunc != nil {
 		return m.allocateFunc(devs)

pkg/kubelet/cm/devicemanager/topology_hints_test.go

@@ -599,6 +599,7 @@ func TestTopologyAlignedAllocation(t *testing.T) {
 			allDevices:            make(map[string]map[string]pluginapi.Device),
 			healthyDevices:        make(map[string]sets.String),
 			allocatedDevices:      make(map[string]sets.String),
+			endpoints:             make(map[string]endpointInfo),
 			podDevices:            make(podDevices),
 			sourcesReady:          &sourcesReadyStub{},
 			activePods:            func() []*v1.Pod { return []*v1.Pod{} },
@@ -607,6 +608,7 @@ func TestTopologyAlignedAllocation(t *testing.T) {
 
 		m.allDevices[tc.resource] = make(map[string]pluginapi.Device)
 		m.healthyDevices[tc.resource] = sets.NewString()
+		m.endpoints[tc.resource] = endpointInfo{}
 
 		for _, d := range tc.devices {
 			m.allDevices[tc.resource][d.ID] = d

test/images/sample-device-plugin/VERSION

@@ -1 +1 @@
-1.1
+1.2

test/images/sample-device-plugin/sampledeviceplugin.go

@@ -86,7 +86,7 @@ func main() {
 	}
 	socketPath := pluginSocksDir + "/dp." + fmt.Sprintf("%d", time.Now().Unix())
 
-	dp1 := dm.NewDevicePluginStub(devs, socketPath, resourceName, false)
+	dp1 := dm.NewDevicePluginStub(devs, socketPath, resourceName, false, false)
 	if err := dp1.Start(); err != nil {
 		panic(err)