From c0ca5091a10417c34192da4d3c064a0fed2a7fdb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Endika=20Guti=C3=A9rrez?=
 <endika.gutierrez@liberehospitality.com>
Date: Tue, 7 Nov 2023 22:24:51 +0100
Subject: [PATCH] Omittable can now be serialized as json (#2839)

---
 graphql/omittable.go      | 23 ++++++++++
 graphql/omittable_test.go | 88 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 111 insertions(+)
 create mode 100644 graphql/omittable_test.go

diff --git a/graphql/omittable.go b/graphql/omittable.go
index 3d1e327904..89716400b6 100644
--- a/graphql/omittable.go
+++ b/graphql/omittable.go
@@ -1,5 +1,7 @@
 package graphql
 
+import "encoding/json"
+
 // Omittable is a wrapper around a value that also stores whether it is set
 // or not.
 type Omittable[T any] struct {
@@ -7,6 +9,11 @@ type Omittable[T any] struct {
 	set   bool
 }
 
+var (
+	_ json.Marshaler   = Omittable[struct{}]{}
+	_ json.Unmarshaler = (*Omittable[struct{}])(nil)
+)
+
 func OmittableOf[T any](value T) Omittable[T] {
 	return Omittable[T]{
 		value: value,
@@ -33,3 +40,19 @@ func (o Omittable[T]) ValueOK() (T, bool) {
 func (o Omittable[T]) IsSet() bool {
 	return o.set
 }
+
+func (o Omittable[T]) MarshalJSON() ([]byte, error) {
+	if !o.set {
+		return []byte("null"), nil
+	}
+	return json.Marshal(o.value)
+}
+
+func (o *Omittable[T]) UnmarshalJSON(bytes []byte) error {
+	err := json.Unmarshal(bytes, &o.value)
+	if err != nil {
+		return err
+	}
+	o.set = true
+	return nil
+}
diff --git a/graphql/omittable_test.go b/graphql/omittable_test.go
new file mode 100644
index 0000000000..79c59178a9
--- /dev/null
+++ b/graphql/omittable_test.go
@@ -0,0 +1,88 @@
+package graphql
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestOmittable_MarshalJSON(t *testing.T) {
+	s := "test"
+	testCases := []struct {
+		name         string
+		input        any
+		expectedJSON string
+	}{
+		{
+			name:         "simple string",
+			input:        struct{ Value Omittable[string] }{Value: OmittableOf("simple string")},
+			expectedJSON: `{"Value": "simple string"}`,
+		},
+		{
+			name:         "string pointer",
+			input:        struct{ Value Omittable[*string] }{Value: OmittableOf(&s)},
+			expectedJSON: `{"Value": "test"}`,
+		},
+		{
+			name:         "omitted integer",
+			input:        struct{ Value Omittable[int] }{},
+			expectedJSON: `{"Value": null}`,
+		},
+		{
+			name:         "omittable omittable",
+			input:        struct{ Value Omittable[Omittable[uint64]] }{Value: OmittableOf(OmittableOf(uint64(42)))},
+			expectedJSON: `{"Value": 42}`,
+		},
+		{
+			name: "omittable struct",
+			input: struct {
+				Value Omittable[struct{ Inner string }]
+			}{Value: OmittableOf(struct{ Inner string }{})},
+			expectedJSON: `{"Value": {"Inner": ""}}`,
+		},
+		{
+			name: "omitted struct",
+			input: struct {
+				Value Omittable[struct{ Inner string }]
+			}{},
+			expectedJSON: `{"Value": null}`,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			data, err := json.Marshal(tc.input)
+			require.NoError(t, err)
+			assert.JSONEq(t, tc.expectedJSON, string(data))
+		})
+	}
+}
+
+func TestOmittable_UnmarshalJSON(t *testing.T) {
+	var s struct {
+		String        Omittable[string]
+		OmittedString Omittable[string]
+		StringPointer Omittable[*string]
+		NullInt       Omittable[int]
+	}
+
+	err := json.Unmarshal([]byte(`
+	{
+		"String": "simple string",
+		"StringPointer": "string pointer",
+		"NullInt": null
+	}`), &s)
+
+	require.NoError(t, err)
+	assert.Equal(t, "simple string", s.String.Value())
+	assert.True(t, s.String.IsSet())
+	assert.False(t, s.OmittedString.IsSet())
+	assert.True(t, s.StringPointer.IsSet())
+	if assert.NotNil(t, s.StringPointer.Value()) {
+		assert.EqualValues(t, "string pointer", *s.StringPointer.Value())
+	}
+	assert.True(t, s.NullInt.IsSet())
+	assert.Zero(t, s.NullInt.Value())
+}