From 49ba70744e7f02ccbd27bece107751d4d569d2bd Mon Sep 17 00:00:00 2001
From: Dmitry Rozhkov <dmitry.rozhkov@intel.com>
Date: Wed, 30 Dec 2020 12:09:00 +0200
Subject: [PATCH] api: improve code coverage

---
 pkg/deviceplugin/manager_test.go |  4 ++++
 pkg/deviceplugin/server_test.go  | 34 ++++++++++++++++++++++++++++----
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/pkg/deviceplugin/manager_test.go b/pkg/deviceplugin/manager_test.go
index a31494a6..16a9ed94 100644
--- a/pkg/deviceplugin/manager_test.go
+++ b/pkg/deviceplugin/manager_test.go
@@ -164,6 +164,10 @@ func (*devicePluginStub) PreStartContainer(*pluginapi.PreStartContainerRequest)
 	return nil
 }
 
+func (*devicePluginStub) GetPreferredAllocation(*pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) {
+	return nil, nil
+}
+
 func TestHandleUpdate(t *testing.T) {
 	tcases := []struct {
 		name            string
diff --git a/pkg/deviceplugin/server_test.go b/pkg/deviceplugin/server_test.go
index 627b321f..56cc8db3 100644
--- a/pkg/deviceplugin/server_test.go
+++ b/pkg/deviceplugin/server_test.go
@@ -492,7 +492,7 @@ func TestGetDevicePluginOptions(t *testing.T) {
 func TestPreStartContainer(t *testing.T) {
 	tcases := []struct {
 		name              string
-		preStartContainer func(*pluginapi.PreStartContainerRequest) error
+		preStartContainer preStartContainerFunc
 		expectedError     bool
 	}{
 		{
@@ -523,9 +523,35 @@ func TestPreStartContainer(t *testing.T) {
 }
 
 func TestGetPreferredAllocation(t *testing.T) {
-	srv := &server{}
-	if _, err := srv.GetPreferredAllocation(context.Background(), nil); err == nil {
-		t.Error("expected an unimplemented error, but got success")
+	tcases := []struct {
+		name                   string
+		getPreferredAllocation getPreferredAllocationFunc
+		expectedError          bool
+	}{
+		{
+			name: "success",
+			getPreferredAllocation: func(*pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) {
+				return nil, nil
+			},
+		},
+		{
+			name:          "error",
+			expectedError: true,
+		},
+	}
+	for _, tc := range tcases {
+		t.Run(tc.name, func(t *testing.T) {
+			srv := &server{
+				getPreferredAllocation: tc.getPreferredAllocation,
+			}
+			_, err := srv.GetPreferredAllocation(context.Background(), nil)
+			if !tc.expectedError && err != nil {
+				t.Errorf("unexpected error: %v", err)
+			} else if tc.expectedError && err == nil {
+				t.Error("didn't failed when expected to fail")
+				return
+			}
+		})
 	}
 }