Skip to content

Commit cbd405d

Browse files
klueskaklueska
committed
Update existing tests in support of GetPreferredallocation()
1 parent a780ccf commit cbd405d

File tree

5 files changed

+68
-56
lines changed

5 files changed

+68
-56
lines changed

‎pkg/kubelet/cm/devicemanager/endpoint_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ func TestAllocate(t *testing.T) {
160160
}
161161

162162
func esetup(t *testing.T, devs []*pluginapi.Device, socket, resourceName string, callback monitorCallback) (*Stub, *endpointImpl) {
163-
p := NewDevicePluginStub(devs, socket, resourceName, false)
163+
p := NewDevicePluginStub(devs, socket, resourceName, false, false)
164164

165165
err := p.Start()
166166
require.NoError(t, err)

‎pkg/kubelet/cm/devicemanager/manager_test.go‎

Lines changed: 63 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -103,55 +103,57 @@ func TestDevicePluginReRegistration(t *testing.T) {
103103
{ID: "Dev3", Health: pluginapi.Healthy},
104104
}
105105
for _, preStartContainerFlag := range []bool{false, true} {
106-
m, ch, p1 := setup(t, devs, nil, socketName, pluginSocketName)
107-
p1.Register(socketName, testResourceName, "")
108-
109-
select {
110-
case <-ch:
111-
case <-time.After(5 * time.Second):
112-
t.Fatalf("timeout while waiting for manager update")
113-
}
114-
capacity, allocatable, _ := m.GetCapacity()
115-
resourceCapacity := capacity[v1.ResourceName(testResourceName)]
116-
resourceAllocatable := allocatable[v1.ResourceName(testResourceName)]
117-
require.Equal(t, resourceCapacity.Value(), resourceAllocatable.Value(), "capacity should equal to allocatable")
118-
require.Equal(t, int64(2), resourceAllocatable.Value(), "Devices are not updated.")
119-
120-
p2 := NewDevicePluginStub(devs, pluginSocketName+".new", testResourceName, preStartContainerFlag)
121-
err = p2.Start()
122-
require.NoError(t, err)
123-
p2.Register(socketName, testResourceName, "")
124-
125-
select {
126-
case <-ch:
127-
case <-time.After(5 * time.Second):
128-
t.Fatalf("timeout while waiting for manager update")
129-
}
130-
capacity, allocatable, _ = m.GetCapacity()
131-
resourceCapacity = capacity[v1.ResourceName(testResourceName)]
132-
resourceAllocatable = allocatable[v1.ResourceName(testResourceName)]
133-
require.Equal(t, resourceCapacity.Value(), resourceAllocatable.Value(), "capacity should equal to allocatable")
134-
require.Equal(t, int64(2), resourceAllocatable.Value(), "Devices shouldn't change.")
135-
136-
// Test the scenario that a plugin re-registers with different devices.
137-
p3 := NewDevicePluginStub(devsForRegistration, pluginSocketName+".third", testResourceName, preStartContainerFlag)
138-
err = p3.Start()
139-
require.NoError(t, err)
140-
p3.Register(socketName, testResourceName, "")
141-
142-
select {
143-
case <-ch:
144-
case <-time.After(5 * time.Second):
145-
t.Fatalf("timeout while waiting for manager update")
106+
for _, getPreferredAllocationFlag := range []bool{false, true} {
107+
m, ch, p1 := setup(t, devs, nil, socketName, pluginSocketName)
108+
p1.Register(socketName, testResourceName, "")
109+
110+
select {
111+
case <-ch:
112+
case <-time.After(5 * time.Second):
113+
t.Fatalf("timeout while waiting for manager update")
114+
}
115+
capacity, allocatable, _ := m.GetCapacity()
116+
resourceCapacity := capacity[v1.ResourceName(testResourceName)]
117+
resourceAllocatable := allocatable[v1.ResourceName(testResourceName)]
118+
require.Equal(t, resourceCapacity.Value(), resourceAllocatable.Value(), "capacity should equal to allocatable")
119+
require.Equal(t, int64(2), resourceAllocatable.Value(), "Devices are not updated.")
120+
121+
p2 := NewDevicePluginStub(devs, pluginSocketName+".new", testResourceName, preStartContainerFlag, getPreferredAllocationFlag)
122+
err = p2.Start()
123+
require.NoError(t, err)
124+
p2.Register(socketName, testResourceName, "")
125+
126+
select {
127+
case <-ch:
128+
case <-time.After(5 * time.Second):
129+
t.Fatalf("timeout while waiting for manager update")
130+
}
131+
capacity, allocatable, _ = m.GetCapacity()
132+
resourceCapacity = capacity[v1.ResourceName(testResourceName)]
133+
resourceAllocatable = allocatable[v1.ResourceName(testResourceName)]
134+
require.Equal(t, resourceCapacity.Value(), resourceAllocatable.Value(), "capacity should equal to allocatable")
135+
require.Equal(t, int64(2), resourceAllocatable.Value(), "Devices shouldn't change.")
136+
137+
// Test the scenario that a plugin re-registers with different devices.
138+
p3 := NewDevicePluginStub(devsForRegistration, pluginSocketName+".third", testResourceName, preStartContainerFlag, getPreferredAllocationFlag)
139+
err = p3.Start()
140+
require.NoError(t, err)
141+
p3.Register(socketName, testResourceName, "")
142+
143+
select {
144+
case <-ch:
145+
case <-time.After(5 * time.Second):
146+
t.Fatalf("timeout while waiting for manager update")
147+
}
148+
capacity, allocatable, _ = m.GetCapacity()
149+
resourceCapacity = capacity[v1.ResourceName(testResourceName)]
150+
resourceAllocatable = allocatable[v1.ResourceName(testResourceName)]
151+
require.Equal(t, resourceCapacity.Value(), resourceAllocatable.Value(), "capacity should equal to allocatable")
152+
require.Equal(t, int64(1), resourceAllocatable.Value(), "Devices of plugin previously registered should be removed.")
153+
p2.Stop()
154+
p3.Stop()
155+
cleanup(t, m, p1)
146156
}
147-
capacity, allocatable, _ = m.GetCapacity()
148-
resourceCapacity = capacity[v1.ResourceName(testResourceName)]
149-
resourceAllocatable = allocatable[v1.ResourceName(testResourceName)]
150-
require.Equal(t, resourceCapacity.Value(), resourceAllocatable.Value(), "capacity should equal to allocatable")
151-
require.Equal(t, int64(1), resourceAllocatable.Value(), "Devices of plugin previously registered should be removed.")
152-
p2.Stop()
153-
p3.Stop()
154-
cleanup(t, m, p1)
155157
}
156158
}
157159

@@ -186,7 +188,7 @@ func TestDevicePluginReRegistrationProbeMode(t *testing.T) {
186188
require.Equal(t, resourceCapacity.Value(), resourceAllocatable.Value(), "capacity should equal to allocatable")
187189
require.Equal(t, int64(2), resourceAllocatable.Value(), "Devices are not updated.")
188190

189-
p2 := NewDevicePluginStub(devs, pluginSocketName+".new", testResourceName, false)
191+
p2 := NewDevicePluginStub(devs, pluginSocketName+".new", testResourceName, false, false)
190192
err = p2.Start()
191193
require.NoError(t, err)
192194
// Wait for the second callback to be issued.
@@ -203,7 +205,7 @@ func TestDevicePluginReRegistrationProbeMode(t *testing.T) {
203205
require.Equal(t, int64(2), resourceAllocatable.Value(), "Devices are not updated.")
204206

205207
// Test the scenario that a plugin re-registers with different devices.
206-
p3 := NewDevicePluginStub(devsForRegistration, pluginSocketName+".third", testResourceName, false)
208+
p3 := NewDevicePluginStub(devsForRegistration, pluginSocketName+".third", testResourceName, false, false)
207209
err = p3.Start()
208210
require.NoError(t, err)
209211
// Wait for the third callback to be issued.
@@ -249,7 +251,7 @@ func setupDeviceManager(t *testing.T, devs []*pluginapi.Device, callback monitor
249251
}
250252

251253
func setupDevicePlugin(t *testing.T, devs []*pluginapi.Device, pluginSocketName string) *Stub {
252-
p := NewDevicePluginStub(devs, pluginSocketName, testResourceName, false)
254+
p := NewDevicePluginStub(devs, pluginSocketName, testResourceName, false, false)
253255
err := p.Start()
254256
require.NoError(t, err)
255257
return p
@@ -549,8 +551,9 @@ func (a *activePodsStub) updateActivePods(newPods []*v1.Pod) {
549551
}
550552

551553
type MockEndpoint struct {
552-
allocateFunc func(devs []string) (*pluginapi.AllocateResponse, error)
553-
initChan chan []string
554+
getPreferredAllocationFunc func(available, mustInclude []string, size int) (*pluginapi.PreferredAllocationResponse, error)
555+
allocateFunc func(devs []string) (*pluginapi.AllocateResponse, error)
556+
initChan chan []string
554557
}
555558

556559
func (m *MockEndpoint) stop() {}
@@ -563,6 +566,13 @@ func (m *MockEndpoint) preStartContainer(devs []string) (*pluginapi.PreStartCont
563566
return &pluginapi.PreStartContainerResponse{}, nil
564567
}
565568

569+
func (m *MockEndpoint) getPreferredAllocation(available, mustInclude []string, size int) (*pluginapi.PreferredAllocationResponse, error) {
570+
if m.getPreferredAllocationFunc != nil {
571+
return m.getPreferredAllocationFunc(available, mustInclude, size)
572+
}
573+
return nil, nil
574+
}
575+
566576
func (m *MockEndpoint) allocate(devs []string) (*pluginapi.AllocateResponse, error) {
567577
if m.allocateFunc != nil {
568578
return m.allocateFunc(devs)

‎pkg/kubelet/cm/devicemanager/topology_hints_test.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,7 @@ func TestTopologyAlignedAllocation(t *testing.T) {
599599
allDevices: make(map[string]map[string]pluginapi.Device),
600600
healthyDevices: make(map[string]sets.String),
601601
allocatedDevices: make(map[string]sets.String),
602+
endpoints: make(map[string]endpointInfo),
602603
podDevices: make(podDevices),
603604
sourcesReady: &sourcesReadyStub{},
604605
activePods: func() []*v1.Pod { return []*v1.Pod{} },
@@ -607,6 +608,7 @@ func TestTopologyAlignedAllocation(t *testing.T) {
607608

608609
m.allDevices[tc.resource] = make(map[string]pluginapi.Device)
609610
m.healthyDevices[tc.resource] = sets.NewString()
611+
m.endpoints[tc.resource] = endpointInfo{}
610612

611613
for _, d := range tc.devices {
612614
m.allDevices[tc.resource][d.ID] = d

‎test/images/sample-device-plugin/VERSION‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
1.1
1+
1.2

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ func main() {
8686
}
8787
socketPath := pluginSocksDir + "/dp." + fmt.Sprintf("%d", time.Now().Unix())
8888

89-
dp1 := dm.NewDevicePluginStub(devs, socketPath, resourceName, false)
89+
dp1 := dm.NewDevicePluginStub(devs, socketPath, resourceName, false, false)
9090
if err := dp1.Start(); err != nil {
9191
panic(err)
9292

0 commit comments

Comments
 (0)