From 1d8389e51320dd8c149a22d631d197526671d413 Mon Sep 17 00:00:00 2001
From: Chris Trombley <chris.trombley@hashicorp.com>
Date: Thu, 4 Jan 2024 15:17:30 -0800
Subject: [PATCH] feat: custom marshaling take 2

---
 models_test.go   | 40 +++++++++++++++++++++++++++++
 response_test.go | 67 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 107 insertions(+)

diff --git a/models_test.go b/models_test.go
index 889142a..a2f1543 100644
--- a/models_test.go
+++ b/models_test.go
@@ -1,7 +1,9 @@
 package jsonapi
 
 import (
+	"encoding/json"
 	"fmt"
+	"reflect"
 	"time"
 )
 
@@ -35,6 +37,44 @@ type TimestampModel struct {
 	RFC3339P *time.Time `jsonapi:"attr,rfc3339p,rfc3339"`
 }
 
+type Unsetable[T any] struct {
+	Value *T
+}
+
+func (t *Unsetable[T]) MarshalJSON() ([]byte, error) {
+	if t == nil {
+		return nil, nil
+	}
+
+	var b []byte
+	var err error
+	if t == nil || t.Value == nil {
+		return json.RawMessage(`null`), nil
+	} else {
+		val := reflect.ValueOf(t.Value)
+		if val.Type().Kind() == reflect.Ptr && val.Elem().Type() == reflect.TypeOf(time.Time{}) {
+			b, err = json.Marshal(val.Elem().Interface().(time.Time).Format(time.RFC3339))
+			if err != nil {
+				return nil, err
+			}
+
+			return b, nil
+		}
+
+		b, err = json.Marshal(t.Value)
+		if err != nil {
+			return nil, err
+		}
+
+		return b, nil
+	}
+}
+
+type CustomTimestampModel struct {
+	ID            int                   `jsonapi:"primary,timestamps"`
+	UnsetableTime *Unsetable[time.Time] `jsonapi:"attr,unsetable"`
+}
+
 type Car struct {
 	ID    *string `jsonapi:"primary,cars"`
 	Make  *string `jsonapi:"attr,make,omitempty"`
diff --git a/response_test.go b/response_test.go
index 599d999..1d3495c 100644
--- a/response_test.go
+++ b/response_test.go
@@ -820,6 +820,73 @@ func TestMarshal_Times(t *testing.T) {
 	}
 }
 
+func TestCustomMarshal_Times(t *testing.T) {
+	aTime := time.Date(2016, 8, 17, 8, 27, 12, 23849, time.UTC)
+
+	for _, tc := range []struct {
+		desc         string
+		input        *CustomTimestampModel
+		verification func(data map[string]interface{}) error
+	}{
+		{
+			desc: "unsetable_nil",
+			input: &CustomTimestampModel{
+				ID:            5,
+				UnsetableTime: nil,
+			},
+			verification: func(root map[string]interface{}) error {
+				v := root["data"].(map[string]interface{})["attributes"].(map[string]interface{})["unsetable"]
+				if got, want := v, (interface{})(nil); got != want {
+					return fmt.Errorf("got %v, want %v", got, want)
+				}
+				return nil
+			},
+		},
+		{
+			desc: "unsetable_value_present",
+			input: &CustomTimestampModel{
+				ID:            5,
+				UnsetableTime: &Unsetable[time.Time]{&aTime},
+			},
+			verification: func(root map[string]interface{}) error {
+				v := root["data"].(map[string]interface{})["attributes"].(map[string]interface{})["unsetable"].(string)
+				if got, want := v, aTime.UTC().Format(time.RFC3339); got != want {
+					return fmt.Errorf("got %v, want %v", got, want)
+				}
+				return nil
+			},
+		},
+		{
+			desc: "unsetable_nil_value",
+			input: &CustomTimestampModel{
+				ID:            5,
+				UnsetableTime: &Unsetable[time.Time]{},
+			},
+			verification: func(root map[string]interface{}) error {
+				v := root["data"].(map[string]interface{})["attributes"].(map[string]interface{})["unsetable"]
+				if got, want := v, (interface{})(nil); got != want {
+					return fmt.Errorf("got %v, want %v", got, want)
+				}
+				return nil
+			},
+		}} {
+		t.Run(tc.desc, func(t *testing.T) {
+			out := bytes.NewBuffer(nil)
+			if err := MarshalPayload(out, tc.input); err != nil {
+				t.Fatal(err)
+			}
+			// Use the standard JSON library to traverse the genereated JSON payload.
+			data := map[string]interface{}{}
+			json.Unmarshal(out.Bytes(), &data)
+			if tc.verification != nil {
+				if err := tc.verification(data); err != nil {
+					t.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
 func TestSupportsLinkable(t *testing.T) {
 	testModel := &Blog{
 		ID:        5,