Add unmarshaling for AssetSetResponse

pkg/kubecost/asset.go

@@ -2031,7 +2031,10 @@ func (n *Node) UnmarshalJSON(b []byte) error {
 		return err
 	}
 
-	n.InterfaceToNode(f)
+	err = n.InterfaceToNode(f)
+	if err != nil {
+		return err
+	}
 
 	return nil
 }
@@ -2622,6 +2625,54 @@ func (sa *SharedAsset) String() string {
 	return toString(sa)
 }
 
+// This type exists because only the assets map of AssetSet is marshaled to
+// json, which makes it impossible to recreate an AssetSet struct. Thus,
+// the type when unmarshaling a marshaled AssetSet,is AssetSetResponse
+type AssetSetResponse struct {
+	assets map[string]Asset
+}
+
+func (asr *AssetSetResponse) UnmarshalJSON(b []byte) error {
+
+	var assetMap map[string]*gojson.RawMessage
+
+	err := gojson.Unmarshal(b, &assetMap)
+	if err != nil {
+		return err
+	}
+
+	newAssetMap := make(map[string]Asset)
+
+	for key, rawMessage := range assetMap {
+
+		var f interface{}
+
+		err := json.Unmarshal(*rawMessage, &f)
+		if err != nil {
+			return err
+		}
+
+		fmap := f.(map[string]interface{})
+
+		switch t := fmap["type"]; t {
+		case "Node":
+
+			var n Node
+			err := n.InterfaceToNode(f)
+
+			if err != nil {
+				return err
+			}
+
+			newAssetMap[key] = &n
+		}
+	}
+
+	asr.assets = newAssetMap
+
+	return nil
+}
+
 // AssetSet stores a set of Assets, each with a unique name, that share
 // a window. An AssetSet is mutable, so treat it like a threadsafe map.
 type AssetSet struct {
@@ -2919,42 +2970,6 @@ func (as *AssetSet) MarshalJSON() ([]byte, error) {
 	return json.Marshal(as.assets)
 }
 
-func (as *AssetSet) UnmarshalJSON(b []byte) error {
-
-	var assetMap map[string]*gojson.RawMessage
-
-	err := gojson.Unmarshal(b, &assetMap)
-	if err != nil {
-		return err
-	}
-
-	var newAssets map[string]Asset
-
-	for _, rawMessage := range assetMap {
-
-		var f interface{}
-
-		err := json.Unmarshal(*rawMessage, &f)
-		if err != nil {
-			return err
-		}
-		fmt.Println(f)
-		fmap := f.(map[string]interface{})
-
-		switch t := fmap["type"]; t {
-		case "Node":
-			var n Node
-			n.InterfaceToNode(f)
-
-			fmt.Println(n)
-		}
-	}
-
-	fmt.Println(assetMap)
-
-	return nil
-}
-
 func (as *AssetSet) Set(asset Asset, aggregateBy []string) error {
 	if as.IsEmpty() {
 		as.Lock()

pkg/kubecost/asset_unmarshal_test.go

@@ -2,7 +2,6 @@ package kubecost
 
 import (
 	"encoding/json"
-	"fmt"
 	"testing"
 	"time"
 )
@@ -124,18 +123,32 @@ func TestAssetset_Unmarshal(t *testing.T) {
 
 	assetList := []Asset{node1, node2, node3}
 
-	assetset1 := NewAssetSet(s, e, assetList...)
+	assetset := NewAssetSet(s, e, assetList...)
+	bytes, _ := json.Marshal(assetset)
 
-	//fmt.Println(assetset1)
+	var assetSetResponse AssetSetResponse
+	assetUnmarshalResponse := &assetSetResponse
 
-	bytes, _ := json.Marshal(assetset1)
+	err := json.Unmarshal(bytes, assetUnmarshalResponse)
 
-	//print(string(bytes))
+	// Check if unmarshal was successful
+	if err != nil {
+		t.Fatalf("AssetSet Unmarshal: unexpected error: %s", err)
+	}
+
+	// For each asset in unmarshaled AssetSetResponse, check if it is equal to the corresponding AssetSet asset
+	for key, asset := range assetset.assets {
 
-	var thisassetset AssetSet
-	assetset2 := &thisassetset
+		if unmarshaledAsset, exists := assetUnmarshalResponse.assets[key]; exists {
 
-	err := json.Unmarshal(bytes, assetset2)
+			if !asset.Equal(unmarshaledAsset) {
+				t.Fatalf("AssetSet Unmarshal: asset at key '%s' from unmarshaled AssetSetResponse does not match corresponding asset from AssetSet", key)
+			}
+
+		} else {
+			t.Fatalf("AssetSet Unmarshal: key '%s' from marshaled AssetSet does not exist in AssetSetResponse", key)
+		}
+
+	}
 
-	fmt.Println(err)
 }