From bb08fd1135f2207670b1bd906c28e00d1036d6af Mon Sep 17 00:00:00 2001
From: Kevin Klues <kklues@nvidia.com>
Date: Thu, 2 Jul 2020 15:15:59 +0000
Subject: [PATCH] Add a simple endpoint test for GetPreferredAllocation()

More extensive tests that exercise the allocation logic are to follow.
---
 pkg/kubelet/cm/devicemanager/endpoint_test.go | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/pkg/kubelet/cm/devicemanager/endpoint_test.go b/pkg/kubelet/cm/devicemanager/endpoint_test.go
index c0a57f2aff4..6496bbb13f0 100644
--- a/pkg/kubelet/cm/devicemanager/endpoint_test.go
+++ b/pkg/kubelet/cm/devicemanager/endpoint_test.go
@@ -159,6 +159,40 @@ func TestAllocate(t *testing.T) {
 	require.Equal(t, resp, respOut)
 }
 
+func TestGetPreferredAllocation(t *testing.T) {
+	socket := path.Join("/tmp", esocketName)
+	callbackCount := 0
+	callbackChan := make(chan int)
+	p, e := esetup(t, []*pluginapi.Device{}, socket, "mock", func(n string, d []pluginapi.Device) {
+		callbackCount++
+		callbackChan <- callbackCount
+	})
+	defer ecleanup(t, p, e)
+
+	resp := &pluginapi.PreferredAllocationResponse{
+		ContainerResponses: []*pluginapi.ContainerPreferredAllocationResponse{
+			{DeviceIDs: []string{"device0", "device1", "device2"}},
+		},
+	}
+
+	p.SetGetPreferredAllocFunc(func(r *pluginapi.PreferredAllocationRequest, devs map[string]pluginapi.Device) (*pluginapi.PreferredAllocationResponse, error) {
+		return resp, nil
+	})
+
+	go e.run()
+	// Wait for the callback to be issued.
+	select {
+	case <-callbackChan:
+		break
+	case <-time.After(time.Second):
+		t.FailNow()
+	}
+
+	respOut, err := e.getPreferredAllocation([]string{}, []string{}, -1)
+	require.NoError(t, err)
+	require.Equal(t, resp, respOut)
+}
+
 func esetup(t *testing.T, devs []*pluginapi.Device, socket, resourceName string, callback monitorCallback) (*Stub, *endpointImpl) {
 	p := NewDevicePluginStub(devs, socket, resourceName, false, false)