[YUNIKORN-1875] requested changes incorporated

Import ordering changes, error handling and logging for gzipReader close, added testcases for decompress.
apache · Sep 15, 2023 · dd15f83 · dd15f83
1 parent 91714fc
commit dd15f83
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 4 deletions.
diff --git a/pkg/common/utils/utils_test.go b/pkg/common/utils/utils_test.go
@@ -23,7 +23,6 @@ import (
 	"compress/gzip"
 	"encoding/base64"
 	"fmt"
-	"github.com/apache/yunikorn-core/pkg/common/configs"
 	"testing"
 	"time"
 
@@ -33,6 +32,7 @@ import (
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
+	"github.com/apache/yunikorn-core/pkg/common/configs"
 	"github.com/apache/yunikorn-k8shim/pkg/common"
 	"github.com/apache/yunikorn-k8shim/pkg/common/constants"
 	"github.com/apache/yunikorn-k8shim/pkg/conf"

diff --git a/pkg/conf/schedulerconf.go b/pkg/conf/schedulerconf.go
@@ -466,6 +466,7 @@ func Decompress(key string, value []byte) (string, string) {
 	n, err := base64.StdEncoding.Decode(decodedValue, value)
 	if err != nil {
 		log.Log(log.ShimConfig).Error("failed to decode schedulerConfig entry", zap.Error(err))
+		return "", ""
 	}
 	decodedValue = decodedValue[:n]
 	splitKey := strings.Split(key, ".")
@@ -475,10 +476,17 @@ func Decompress(key string, value []byte) (string, string) {
 		gzReader, err := gzip.NewReader(reader)
 		if err != nil {
 			log.Log(log.ShimConfig).Error("failed to decompress decoded schedulerConfig entry", zap.Error(err))
+			return "", ""
 		}
+		defer func() {
+			if err := gzReader.Close(); err != nil {
+				log.Log(log.ShimConfig).Debug("gzip Reader could not be closed ", zap.Error(err))
+			}
+		}()
 		decompressedBytes, err := io.ReadAll(gzReader)
 		if err != nil {
 			log.Log(log.ShimConfig).Error("failed to decompress decoded schedulerConfig entry", zap.Error(err))
+			return "", ""
 		}
 		uncompressedData = string(decompressedBytes)
 	}

diff --git a/pkg/conf/schedulerconf_test.go b/pkg/conf/schedulerconf_test.go
@@ -22,14 +22,15 @@ import (
 	"compress/gzip"
 	"encoding/base64"
 	"fmt"
-	"github.com/apache/yunikorn-core/pkg/common/configs"
 	"reflect"
+	"strings"
 	"testing"
 	"time"
 
 	"gotest.tools/v3/assert"
 	v1 "k8s.io/api/core/v1"
 
+	"github.com/apache/yunikorn-core/pkg/common/configs"
 	"github.com/apache/yunikorn-k8shim/pkg/common/constants"
 )
 
@@ -80,17 +81,35 @@ func TestDecompress(t *testing.T) {
 	var b bytes.Buffer
 	gzWriter := gzip.NewWriter(&b)
 	if _, err := gzWriter.Write([]byte(configs.DefaultSchedulerConfig)); err != nil {
-		t.Fatal("expected nil, got error while compressing test schedulerConfig")
+		assert.NilError(t, err, "expected nil, got error while compressing test schedulerConfig")
 	}
 	if err := gzWriter.Close(); err != nil {
+		assert.NilError(t, err, "expected nil, got error")
 		t.Fatal("expected nil, got error")
 	}
 	encodedConfigString := make([]byte, base64.StdEncoding.EncodedLen(len(b.Bytes())))
 	base64.StdEncoding.Encode(encodedConfigString, b.Bytes())
-	_, decodedConfigString := Decompress("queues.yaml."+constants.GzipSuffix, encodedConfigString)
+	key, decodedConfigString := Decompress("queues.yaml."+constants.GzipSuffix, encodedConfigString)
+	assert.Equal(t, "queues.yaml", key)
 	assert.Equal(t, configs.DefaultSchedulerConfig, decodedConfigString)
 }
 
+func TestDecompressUnkownKey(t *testing.T) {
+	encodedConfigString := make([]byte, base64.StdEncoding.EncodedLen(len([]byte(configs.DefaultSchedulerConfig))))
+	base64.StdEncoding.Encode(encodedConfigString, []byte(configs.DefaultSchedulerConfig))
+	key, decodedConfigString := Decompress("queues.yaml.bin", encodedConfigString)
+	assert.Equal(t, "queues.yaml", key)
+	assert.Assert(t, len(decodedConfigString) == 0, "expected decodedConfigString to be nil")
+}
+
+func TestDecompressBadCompression(t *testing.T) {
+	encodedConfigString := make([]byte, base64.StdEncoding.EncodedLen(len([]byte(configs.DefaultSchedulerConfig))))
+	base64.StdEncoding.Encode(encodedConfigString, []byte(configs.DefaultSchedulerConfig))
+	key, decodedConfigString := Decompress("queues.yaml."+constants.GzipSuffix, encodedConfigString)
+	assert.Equal(t, "", key)
+	assert.Assert(t, !strings.EqualFold(configs.DefaultSchedulerConfig, decodedConfigString), "expected decodedConfigString to be nil")
+}
+
 func TestParseConfigMap(t *testing.T) {
 	testCases := []struct {
 		name  string