kubelet dra: add unit tests for checkpoint

Signed-off-by: Moshe Levi <moshele@nvidia.com>

pkg/kubelet/cm/dra/state/state_checkpoint_test.go

@@ -0,0 +1,151 @@
+/*
+Copyright 2020 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package state
+
+import (
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager"
+	testutil "k8s.io/kubernetes/pkg/kubelet/cm/cpumanager/state/testing"
+)
+
+const testingCheckpoint = "dramanager_checkpoint_test"
+
+// assertStateEqual marks provided test as failed if provided states differ
+func assertStateEqual(t *testing.T, restoredState, expectedState ClaimInfoStateList) {
+	assert.Equal(t, expectedState, restoredState, "expected ClaimInfoState does not equal to restored one")
+}
+
+func TestCheckpointGetOrCreate(t *testing.T) {
+	testCases := []struct {
+		description       string
+		checkpointContent string
+		expectedError     string
+		expectedState     ClaimInfoStateList
+	}{
+		{
+			"Create non-existing checkpoint",
+			"",
+			"",
+			[]ClaimInfoState{},
+		},
+		{
+			"Restore valid checkpoint",
+			`{"version":"v1","entries":[{"DriverName":"test-driver.cdi.k8s.io","ClaimUID":"067798be-454e-4be4-9047-1aa06aea63f7","ClaimName":"example","Namespace":"default","PodUIDs":{"139cdb46-f989-4f17-9561-ca10cfb509a6":{}},"CdiDevices":["example.com/example=cdi-example"]}],"checksum":2939981547}`,
+			"",
+			[]ClaimInfoState{{
+				DriverName: "test-driver.cdi.k8s.io",
+				ClaimUID:   "067798be-454e-4be4-9047-1aa06aea63f7",
+				ClaimName:  "example",
+				Namespace:  "default",
+				PodUIDs:    sets.New("139cdb46-f989-4f17-9561-ca10cfb509a6"),
+				CdiDevices: []string{"example.com/example=cdi-example"},
+			},
+			},
+		},
+		{
+			"Restore checkpoint with invalid checksum",
+			`{"version":"v1","entries":[{"DriverName":"test-driver.cdi.k8s.io","ClaimUID":"067798be-454e-4be4-9047-1aa06aea63f7","ClaimName":"example","Namespace":"default","PodUIDs":{"139cdb46-f989-4f17-9561-ca10cfb509a6":{}},"CdiDevices":["example.com/example=cdi-example"]}],"checksum":2939981548}`,
+			"checkpoint is corrupted",
+			[]ClaimInfoState{},
+		},
+		{
+			"Restore checkpoint with invalid JSON",
+			`{`,
+			"unexpected end of JSON input",
+			[]ClaimInfoState{},
+		},
+	}
+
+	// create temp dir
+	testingDir, err := os.MkdirTemp("", "dramanager_state_test")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(testingDir)
+
+	// create checkpoint manager for testing
+	cpm, err := checkpointmanager.NewCheckpointManager(testingDir)
+	assert.NoError(t, err, "could not create testing checkpoint manager")
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			// ensure there is no previous checkpoint
+			assert.NoError(t, cpm.RemoveCheckpoint(testingCheckpoint), "could not remove testing checkpoint")
+
+			// prepare checkpoint for testing
+			if strings.TrimSpace(tc.checkpointContent) != "" {
+				checkpoint := &testutil.MockCheckpoint{Content: tc.checkpointContent}
+				assert.NoError(t, cpm.CreateCheckpoint(testingCheckpoint, checkpoint), "could not create testing checkpoint")
+			}
+
+			var state ClaimInfoStateList
+
+			checkpointState, err := NewCheckpointState(testingDir, testingCheckpoint)
+			if err == nil {
+				state, err = checkpointState.GetOrCreate()
+			}
+			if strings.TrimSpace(tc.expectedError) != "" {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), tc.expectedError)
+			} else {
+				assert.NoError(t, err, "unexpected error while creating checkpointState")
+				// compare state after restoration with the one expected
+				assertStateEqual(t, state, tc.expectedState)
+			}
+		})
+	}
+}
+
+func TestCheckpointStateStore(t *testing.T) {
+	claimInfoState := ClaimInfoState{
+		DriverName: "test-driver.cdi.k8s.io",
+		ClaimUID:   "067798be-454e-4be4-9047-1aa06aea63f7",
+		ClaimName:  "example",
+		Namespace:  "default",
+		PodUIDs:    sets.New("139cdb46-f989-4f17-9561-ca10cfb509a6"),
+		CdiDevices: []string{"example.com/example=cdi-example"},
+	}
+
+	expectedCheckpoint := `{"version":"v1","entries":[{"DriverName":"test-driver.cdi.k8s.io","ClaimUID":"067798be-454e-4be4-9047-1aa06aea63f7","ClaimName":"example","Namespace":"default","PodUIDs":{"139cdb46-f989-4f17-9561-ca10cfb509a6":{}},"CdiDevices":["example.com/example=cdi-example"]}],"checksum":2939981547}`
+
+	// create temp dir
+	testingDir, err := os.MkdirTemp("", "dramanager_state_test")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(testingDir)
+
+	cpm, err := checkpointmanager.NewCheckpointManager(testingDir)
+	assert.NoError(t, err, "could not create testing checkpoint manager")
+	assert.NoError(t, cpm.RemoveCheckpoint(testingCheckpoint), "could not remove testing checkpoint")
+
+	cs, err := NewCheckpointState(testingDir, testingCheckpoint)
+	assert.NoError(t, err, "could not create testing checkpointState instance")
+	err = cs.Store(ClaimInfoStateList{claimInfoState})
+	assert.NoError(t, err, "could not store ClaimInfoState")
+	checkpoint := NewDRAManagerCheckpoint()
+	cpm.GetCheckpoint(testingCheckpoint, checkpoint)
+	checkpointData, err := checkpoint.MarshalCheckpoint()
+	assert.NoError(t, err, "could not Marshal Checkpoint")
+	assert.Equal(t, expectedCheckpoint, string(checkpointData), "expected ClaimInfoState does not equal to restored one")
+}