test: refactor previous json marshal test

Signed-off-by: DenChenn <[email protected]>
flyteorg · Oct 23, 2024 · 34766cb · 34766cb
1 parent b5f19be
commit 34766cb
Showing 1 changed file with 19 additions and 13 deletions.
diff --git a/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/error_test.go b/flytepropeller/pkg/apis/flyteworkflow/v1alpha1/error_test.go
@@ -1,7 +1,6 @@
 package v1alpha1
 
 import (
-	"encoding/json"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -10,21 +9,28 @@ import (
 )
 
 func TestExecutionErrorJSONMarshalling(t *testing.T) {
-	execError := &core.ExecutionError{
-		Code:     "TestCode",
-		Message:  "Test error message",
-		ErrorUri: "Test error uri",
+	execError := ExecutionError{
+		&core.ExecutionError{
+			Code:     "TestCode",
+			Message:  "Test error message",
+			ErrorUri: "Test error uri",
+		},
 	}
 
-	execErr := &ExecutionError{ExecutionError: execError}
-	data, jErr := json.Marshal(execErr)
-	assert.Nil(t, jErr)
+	expected, mockErr := mockMarshalPbToBytes(execError.ExecutionError)
+	assert.Nil(t, mockErr)
 
-	newExecErr := &ExecutionError{}
-	uErr := json.Unmarshal(data, newExecErr)
+	// MarshalJSON
+	execErrorBytes, mErr := execError.MarshalJSON()
+	assert.Nil(t, mErr)
+	assert.Equal(t, expected, execErrorBytes)
+
+	// UnmarshalJSON
+	execErrorObj := &ExecutionError{}
+	uErr := execErrorObj.UnmarshalJSON(execErrorBytes)
 	assert.Nil(t, uErr)
 
-	assert.Equal(t, execError.Code, newExecErr.ExecutionError.Code)
-	assert.Equal(t, execError.Message, newExecErr.ExecutionError.Message)
-	assert.Equal(t, execError.ErrorUri, newExecErr.ExecutionError.ErrorUri)
+	assert.Equal(t, execError.Code, execErrorObj.Code)
+	assert.Equal(t, execError.Message, execError.Message)
+	assert.Equal(t, execError.ErrorUri, execErrorObj.ErrorUri)
 }