Просмотр исходного кода

Implement UnmarshalJSON for Window type

Signed-off-by: Michael Dresser <michaelmdresser@gmail.com>
Michael Dresser 3 лет назад
Родитель
Сommit
48e60f4d13
2 измененных файлов с 84 добавлено и 5 удалено
  1. 39 4
      pkg/kubecost/window.go
  2. 45 1
      pkg/kubecost/window_test.go

+ 39 - 4
pkg/kubecost/window.go

@@ -2,17 +2,17 @@ package kubecost
 
 import (
 	"bytes"
+	"encoding/json"
 	"fmt"
-	"github.com/opencost/opencost/pkg/log"
 	"math"
 	"regexp"
 	"strconv"
 	"time"
 
-	"github.com/opencost/opencost/pkg/util/timeutil"
-
 	"github.com/opencost/opencost/pkg/env"
+	"github.com/opencost/opencost/pkg/log"
 	"github.com/opencost/opencost/pkg/thanos"
+	"github.com/opencost/opencost/pkg/util/timeutil"
 )
 
 const (
@@ -479,7 +479,6 @@ func (w Window) IsOpen() bool {
 	return w.start == nil || w.end == nil
 }
 
-// TODO:CLEANUP make this unmarshalable (make Start and End public)
 func (w Window) MarshalJSON() ([]byte, error) {
 	buffer := bytes.NewBufferString("{")
 	if w.start != nil {
@@ -496,6 +495,42 @@ func (w Window) MarshalJSON() ([]byte, error) {
 	return buffer.Bytes(), nil
 }
 
+func (w *Window) UnmarshalJSON(bs []byte) error {
+	// Due to the behavior of our custom MarshalJSON, we unmarshal as strings
+	// and then manually handle the weird quoted "null" case.
+	type PubWindow struct {
+		Start string `json:"start"`
+		End   string `json:"end"`
+	}
+	var pw PubWindow
+	err := json.Unmarshal(bs, &pw)
+	if err != nil {
+		return fmt.Errorf("half unmarshal: %w", err)
+	}
+
+	var start *time.Time
+	var end *time.Time
+
+	if pw.Start != "null" {
+		t, err := time.Parse(time.RFC3339, pw.Start)
+		if err != nil {
+			return fmt.Errorf("parsing start as RFC3339: %w", err)
+		}
+		start = &t
+	}
+	if pw.End != "null" {
+		t, err := time.Parse(time.RFC3339, pw.End)
+		if err != nil {
+			return fmt.Errorf("parsing end as RFC3339: %w", err)
+		}
+		end = &t
+	}
+
+	w.start = start
+	w.end = end
+	return nil
+}
+
 func (w Window) Minutes() float64 {
 	if w.IsOpen() {
 		return math.Inf(1)

+ 45 - 1
pkg/kubecost/window_test.go

@@ -1,12 +1,15 @@
 package kubecost
 
 import (
+	"encoding/json"
 	"fmt"
-	"github.com/opencost/opencost/pkg/util/timeutil"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
+	"github.com/opencost/opencost/pkg/util/timeutil"
+
 	"github.com/opencost/opencost/pkg/env"
 )
 
@@ -1146,3 +1149,44 @@ func TestWindow_GetWindowsForQueryWindow(t *testing.T) {
 		})
 	}
 }
+
+func TestMarshalUnmarshal(t *testing.T) {
+	t1 := time.Date(2023, 03, 11, 01, 29, 15, 0, time.UTC)
+	t2 := t1.Add(8 * time.Minute)
+	cases := []struct {
+		w Window
+	}{
+		{
+			w: NewClosedWindow(t1, t2),
+		},
+		{
+			w: NewWindow(&t1, nil),
+		},
+		{
+			w: NewWindow(nil, &t2),
+		},
+		{
+			w: NewWindow(nil, nil),
+		},
+	}
+
+	for _, c := range cases {
+		name := c.w.String()
+		t.Run(name, func(t *testing.T) {
+			marshaled, err := json.Marshal(c.w)
+			if err != nil {
+				t.Fatalf("marshaling: %s", err)
+			}
+
+			var unmarshaledW Window
+			err = json.Unmarshal(marshaled, &unmarshaledW)
+			if err != nil {
+				t.Fatalf("unmarshaling: %s", err)
+			}
+
+			if diff := cmp.Diff(c.w, unmarshaledW); len(diff) > 0 {
+				t.Errorf(diff)
+			}
+		})
+	}
+}