add storage validation helper, more tests

pkg/storage/storage.go

@@ -3,6 +3,8 @@ package storage
 import (
 	"os"
 	"time"
+
+	"github.com/pkg/errors"
 )
 
 // DirDelim is the delimiter used to model a directory structure in an object store bucket.
@@ -51,6 +53,41 @@ type Storage interface {
 	List(path string) ([]*StorageInfo, error)
 }
 
+// Validate uses the provided storage implementation to write a test file to the store, followed by a removal.
+func Validate(storage Storage) error {
+	const testPath = "tmp/test.txt"
+	const testContent = "test"
+
+	// attempt to read a path
+	_, err := storage.Exists(testPath)
+	if err != nil {
+		return errors.Wrap(err, "Failed to check if path exists")
+	}
+
+	// attempt to write a path
+	err = storage.Write(testPath, []byte(testContent))
+	if err != nil {
+		return errors.Wrap(err, "Failed to write data to storage")
+	}
+
+	// attempt to read the path
+	data, err := storage.Read(testPath)
+	if err != nil {
+		return errors.Wrap(err, "Failed to read data from storage")
+	}
+	if string(data) != testContent {
+		return errors.New("Failed to read the expected data from storage")
+	}
+
+	// delete the path
+	err = storage.Remove(testPath)
+	if err != nil {
+		return errors.Wrap(err, "Failed to remove data from storage")
+	}
+
+	return nil
+}
+
 // IsNotExist returns true if the error provided from a storage object is DoesNotExist
 func IsNotExist(err error) bool {
 	if err == nil {

pkg/storage/storagetypes.go

@@ -1,6 +1,10 @@
 package storage
 
-import "strings"
+import (
+	"strings"
+
+	"github.com/kubecost/cost-model/pkg/util/json"
+)
 
 /*
  NOTE: This format is to provide monitoring a simple way to identify the storage type with
@@ -20,6 +24,37 @@ const (
 	StorageTypeBucketAzure StorageType = "bucket|azure"
 )
 
+// jsonIR is a json intermediate representation of a StorageType
+type jsonIR struct {
+	BackendType  string `json:"backendType"`
+	ProviderType string `json:"providerType,omitempty"`
+}
+
+// MarshalJSON implements the json.Marshaler interface for encoding a StorageType.
+func (st StorageType) MarshalJSON() ([]byte, error) {
+	return json.Marshal(jsonIR{
+		BackendType:  st.BackendType(),
+		ProviderType: st.ProviderType(),
+	})
+}
+
+// UnmarshalJSON implements the json.Unmarshaler interface for decoding a StorageType.
+func (st *StorageType) UnmarshalJSON(data []byte) error {
+	var ir jsonIR
+	err := json.Unmarshal(data, &ir)
+	if err != nil {
+		return err
+	}
+
+	str := ir.BackendType
+	if ir.ProviderType != "" {
+		str += "|" + ir.ProviderType
+	}
+
+	*st = StorageType(str)
+	return nil
+}
+
 // IsFileStorage returns true if the StorageType is a file storage type.
 func (st StorageType) IsFileStorage() bool {
 	return st.BackendType() == "file"

pkg/storage/storagetypes_test.go

@@ -1,6 +1,10 @@
 package storage
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/kubecost/cost-model/pkg/util/json"
+)
 
 func assert(t *testing.T, condition bool, msg string) {
 	if !condition {
@@ -8,6 +12,12 @@ func assert(t *testing.T, condition bool, msg string) {
 	}
 }
 
+func assertEq[T comparable](t *testing.T, got, expected T) {
+	if got != expected {
+		t.Errorf("Failed Equality Assertion:\n  Got: %v\n  Exp: %v", got, expected)
+	}
+}
+
 func TestStorageTypes(t *testing.T) {
 	fileType := StorageTypeFile
 	s3Type := StorageTypeBucketS3
@@ -28,5 +38,144 @@ func TestStorageTypes(t *testing.T) {
 	assert(t, s3Type.IsBucketStorage(), "StorageTypeBucketS3.IsBucketStorage() should return true")
 	assert(t, gcsType.IsBucketStorage(), "StorageTypeBucketGCS.IsBucketStorage() should return true")
 	assert(t, azureType.IsBucketStorage(), "StorageTypeBucketAzure.IsBucketStorage() should return true")
+}
+
+func TestJSONEncodeStorageType(t *testing.T) {
+	fileType := StorageTypeFile
+	s3Type := StorageTypeBucketS3
+	gcsType := StorageTypeBucketGCS
+	azureType := StorageTypeBucketAzure
+
+	var data []byte
+	var err error
+
+	data, err = json.Marshal(fileType)
+	if err != nil {
+		t.Error(err)
+	}
+	assert(t, string(data) == `{"backendType":"file"}`, "json.Marshal(StorageTypeFile) should return '\"file\"'")
+
+	data, err = json.Marshal(s3Type)
+	if err != nil {
+		t.Error(err)
+	}
+	assert(t, string(data) == `{"backendType":"bucket","providerType":"s3"}`, "json.Marshal(StorageTypeBucketS3) should return '\"bucket|s3\"'")
+
+	data, err = json.Marshal(gcsType)
+	if err != nil {
+		t.Error(err)
+	}
+	assert(t, string(data) == `{"backendType":"bucket","providerType":"gcs"}`, "json.Marshal(StorageTypeBucketGCS) should return '\"bucket|gcs\"'")
+
+	data, err = json.Marshal(azureType)
+	if err != nil {
+		t.Error(err)
+	}
+	assert(t, string(data) == `{"backendType":"bucket","providerType":"azure"}`, "json.Marshal(StorageTypeBucketAzure) should return '\"bucket|azure\"'")
+}
+
+func TestJSONDecodeStorageType(t *testing.T) {
+	fileType := StorageTypeFile
+	s3Type := StorageTypeBucketS3
+	gcsType := StorageTypeBucketGCS
+	azureType := StorageTypeBucketAzure
+
+	var st StorageType
+
+	data := []byte(`{"backendType":"file"}`)
+	err := json.Unmarshal(data, &st)
+	if err != nil {
+		t.Error(err)
+	}
+	assert(t, st == fileType, "json.Unmarshal() should return StorageTypeFile")
 
+	data = []byte(`{"backendType":"bucket","providerType":"s3"}`)
+	err = json.Unmarshal(data, &st)
+	if err != nil {
+		t.Error(err)
+	}
+	assert(t, st == s3Type, "json.Unmarshal() should return StorageTypeBucketS3")
+
+	data = []byte(`{"backendType":"bucket","providerType":"gcs"}`)
+	err = json.Unmarshal(data, &st)
+	if err != nil {
+		t.Error(err)
+	}
+	assert(t, st == gcsType, "json.Unmarshal() should return StorageTypeBucketGCS")
+
+	data = []byte(`{"backendType":"bucket","providerType":"azure"}`)
+	err = json.Unmarshal(data, &st)
+	if err != nil {
+		t.Error(err)
+	}
+	assert(t, st == azureType, "json.Unmarshal() should return StorageTypeBucketAzure")
+}
+
+type TestWrapper struct {
+	Foo         string      `json:"foo"`
+	Prop        int         `json:"prop"`
+	StorageType StorageType `json:"storageType"`
+}
+
+func TestJSONEncodeStorageTypeWrapped(t *testing.T) {
+	tw := TestWrapper{
+		Foo:         "bar",
+		Prop:        42,
+		StorageType: StorageTypeFile,
+	}
+
+	var data []byte
+	var err error
+
+	data, err = json.Marshal(tw)
+	if err != nil {
+		t.Error(err)
+	}
+	assertEq(t, string(data), `{"foo":"bar","prop":42,"storageType":{"backendType":"file"}}`)
+
+	tw = TestWrapper{
+		Foo:         "bar",
+		Prop:        42,
+		StorageType: StorageTypeBucketS3,
+	}
+
+	data, err = json.Marshal(tw)
+	if err != nil {
+		t.Error(err)
+	}
+	assertEq(t, string(data), `{"foo":"bar","prop":42,"storageType":{"backendType":"bucket","providerType":"s3"}}`)
+}
+
+func TestJSONDecodeStorageTypeWrapped(t *testing.T) {
+	tw := TestWrapper{
+		Foo:         "bar",
+		Prop:        42,
+		StorageType: StorageTypeFile,
+	}
+
+	var stw TestWrapper
+
+	data := []byte(`{"foo":"bar","prop":42,"storageType":{"backendType":"file"}}`)
+	err := json.Unmarshal(data, &stw)
+	if err != nil {
+		t.Error(err)
+	}
+	assertEq(t, stw.Foo, tw.Foo)
+	assertEq(t, stw.Prop, tw.Prop)
+	assertEq(t, stw.StorageType, tw.StorageType)
+
+	tw = TestWrapper{
+		Foo:         "bar",
+		Prop:        42,
+		StorageType: StorageTypeBucketS3,
+	}
+
+	data = []byte(`{"foo":"bar","prop":42,"storageType":{"backendType":"bucket","providerType":"s3"}}`)
+	err = json.Unmarshal(data, &stw)
+	if err != nil {
+		t.Error(err)
+	}
+	assertEq(t, stw.Foo, tw.Foo)
+	assertEq(t, stw.Prop, tw.Prop)
+	assertEq(t, stw.StorageType, tw.StorageType)
 }