From 87432c04b501e78b7f3caf030a1cfe70747ca935 Mon Sep 17 00:00:00 2001
From: Fahad Siddiqui <6371229+fahadsiddiqui@users.noreply.github.com>
Date: Tue, 3 Sep 2024 14:01:29 +0500
Subject: [PATCH] Update NullString omitempty Logic

---
 nulltypes_test.go | 52 +++++++++++++++++++++++++++++++++++------------
 string.go         | 33 +++++++++++++++++-------------
 2 files changed, 58 insertions(+), 27 deletions(-)

diff --git a/nulltypes_test.go b/nulltypes_test.go
index d715ab1..93526d0 100644
--- a/nulltypes_test.go
+++ b/nulltypes_test.go
@@ -7,7 +7,7 @@ import (
 )
 
 type TestStruct struct {
-	Name          NullString
+	Name          NullString `json:"Name,omitempty"`
 	IsMarried     NullBool
 	Height        NullFloat64
 	Age           NullInt32
@@ -24,6 +24,8 @@ type TestStruct struct {
 var tz, _ = time.LoadLocation("UTC")
 var timeToday time.Time = time.Date(2020, 10, 16, 01, 04, 30, 0, tz)
 var tStr = `{"Name":"John Doe","IsMarried":true,"Height":170.4,"Age":40,"Income":4000000,"BornAt":"2020-10-16T01:04:30Z","NameNull":null,"IsMarriedNull":null,"HeightNull":null,"AgeNull":null,"IncomeNull":null,"BornAtNull":null}`
+var tStrWOmitEmpty = `{"Name":null,"IsMarried":true,"Height":170.4,"Age":40,"Income":4000000,"BornAt":"2020-10-16T01:04:30Z","NameNull":null,"IsMarriedNull":null,"HeightNull":null,"AgeNull":null,"IncomeNull":null,"BornAtNull":null}`
+
 var tObj = TestStruct{
 	Name:      NullString{String: "John Doe", Valid: true},
 	IsMarried: NullBool{Bool: true, Valid: true},
@@ -33,28 +35,52 @@ var tObj = TestStruct{
 	BornAt:    NullTime{Time: timeToday, Valid: true},
 }
 
-func TestMarshalJSON(t *testing.T) {
+var tObjWithEmptyString = TestStruct{
+	IsMarried: NullBool{Bool: true, Valid: true},
+	Height:    NullFloat64{Float64: 170.4, Valid: true},
+	Age:       NullInt32{Int32: 40, Valid: true},
+	Income:    NullInt64{Int64: 4000000, Valid: true},
+	BornAt:    NullTime{Time: timeToday, Valid: true},
+}
+
+func TestUnmarshalJSON(t *testing.T) {
+	// Test with all fields present
 	var ts1 TestStruct
-	b := []byte(tStr)
-	err := json.Unmarshal(b, &ts1)
+	err := json.Unmarshal([]byte(tStr), &ts1)
 	if err != nil {
-		t.Errorf("Something went bad while unmarshaling %+v obj...", tObj)
+		t.Errorf("Error unmarshaling JSON: %v", err)
 	}
-
-	t.Logf("%+v", ts1)
-
 	if ts1 != tObj {
-		t.Errorf("Unmarshalling string didn't work as expected, \n%+v\n!=\n%+v\n, somehow!", ts1, tObj)
+		t.Errorf("Unmarshaling didn't work as expected:\nGot: %+v\nWant: %+v", ts1, tObj)
+	}
+
+	// Test with Name field omitted (Name.Valid should be false)
+	var ts2 TestStruct
+	err = json.Unmarshal([]byte(tStrWOmitEmpty), &ts2)
+	if err != nil {
+		t.Errorf("Error unmarshaling JSON: %v", err)
+	}
+	if ts2 != tObjWithEmptyString {
+		t.Errorf("Unmarshaling didn't work as expected:\nGot: %+v\nWant: %+v", ts2, tObjWithEmptyString)
 	}
 }
 
-func TestUnmarshalJSON(t *testing.T) {
+func TestMarshalJSON(t *testing.T) {
+	// Test marshaling with all fields present
 	b, err := json.Marshal(tObj)
 	if err != nil {
-		t.Errorf("Something went bad while marshaling %+v obj...", tObj)
+		t.Errorf("Error marshaling object: %v", err)
 	}
-
 	if string(b) != tStr {
-		t.Errorf("Marshalling string didn't work as expected, %v != %v, somehow!", string(b), tStr)
+		t.Errorf("Marshaling didn't work as expected:\nGot: %v\nWant: %v", string(b), tStr)
+	}
+
+	// Test marshaling with Name field omitted (Name.Valid is false)
+	b, err = json.Marshal(tObjWithEmptyString)
+	if err != nil {
+		t.Errorf("Error marshaling object: %v", err)
+	}
+	if string(b) != tStrWOmitEmpty {
+		t.Errorf("Marshaling didn't work as expected:\nGot: %v\nWant: %v", string(b), tStrWOmitEmpty)
 	}
 }
diff --git a/string.go b/string.go
index 6840376..90e88de 100644
--- a/string.go
+++ b/string.go
@@ -3,6 +3,7 @@ package nulltypes
 import (
 	"database/sql/driver"
 	"encoding/json"
+	"fmt"
 )
 
 // NullString is a wrapper around string
@@ -23,7 +24,8 @@ func String(s string) NullString {
 // whenever it is of type NullString
 func (ns NullString) MarshalJSON() ([]byte, error) {
 	if !ns.Valid {
-		return json.Marshal(nil)
+		// Return nil to ensure the field is omitted
+		return []byte("null"), nil
 	}
 	return json.Marshal(ns.String)
 }
@@ -31,31 +33,34 @@ func (ns NullString) MarshalJSON() ([]byte, error) {
 // UnmarshalJSON method is called by json.Unmarshal,
 // whenever it is of type NullString
 func (ns *NullString) UnmarshalJSON(b []byte) error {
-	var s *string
+	if string(b) == "null" {
+		ns.Valid = false
+		ns.String = ""
+		return nil
+	}
+	var s string
 	if err := json.Unmarshal(b, &s); err != nil {
 		return err
 	}
-	if s != nil {
-		ns.Valid = true
-		ns.String = *s
-	} else {
-		ns.Valid = false
-	}
+	ns.Valid = true
+	ns.String = s
 	return nil
 }
 
-// Scan satisfies the sql.scanner interface
+// Scan satisfies the sql.Scanner interface
 func (ns *NullString) Scan(value interface{}) error {
-	rt, ok := value.(string)
-	if ok {
-		*ns = NullString{rt, true}
-	} else {
+	switch v := value.(type) {
+	case string:
+		*ns = NullString{v, true}
+	case nil:
 		*ns = NullString{"", false}
+	default:
+		return fmt.Errorf("unable to scan type %T into NullString", value)
 	}
 	return nil
 }
 
-// Value satisfies the driver.Value interface
+// Value satisfies the driver.Valuer interface
 func (ns NullString) Value() (driver.Value, error) {
 	if ns.Valid {
 		return ns.String, nil