feat: customize workload deserialization and serialization by type (#456

)

#### What type of PR is this?
/kind feature

#### What this PR does / why we need it:
Customize workload deserialization and serialization by type.

pkg/models/appconfiguration/workload/common.go

@@ -7,11 +7,11 @@ import "kusionstack.io/kusion/pkg/models/appconfiguration/workload/container"
 // these common attributes.
 type WorkloadBase struct {
 	// The templates of containers to be ran.
-	Containers map[string]container.Container `yaml:"containers" json:"containers"`
+	Containers map[string]container.Container `yaml:"containers,omitempty" json:"containers,omitempty"`
 
 	// The number of containers that should be ran.
 	// Default is 2 to meet high availability requirements.
-	Replicas int `yaml:"replicas" json:"replicas"`
+	Replicas int `yaml:"replicas,omitempty" json:"replicas,omitempty"`
 
 	// Labels and annotations can be used to attach arbitrary metadata
 	// as key-value pairs to resources.

pkg/models/appconfiguration/workload/job.go

@@ -4,9 +4,9 @@ package workload
 // application code. This is typically used for tasks that take from a
 // few senconds to a few days to complete.
 type Job struct {
-	*WorkloadBase `yaml:",inline" json:",inline"`
+	WorkloadBase `yaml:",inline" json:",inline"`
 
 	// The scheduling strategy in Cron format.
 	// More info: https://en.wikipedia.org/wiki/Cron.
-	Schedule string `yaml:"schedule" json:"schedule"`
+	Schedule string `yaml:"schedule,omitempty" json:"schedule,omitempty"`
 }

pkg/models/appconfiguration/workload/service.go

@@ -5,5 +5,5 @@ package workload
 // applications that should "never" go down, and handle short-lived
 // latency-sensitive web requests, or events.
 type Service struct {
-	*WorkloadBase `yaml:",inline" json:",inline"`
+	WorkloadBase `yaml:",inline" json:",inline"`
 }

pkg/models/appconfiguration/workload/workload.go

@@ -1,14 +1,70 @@
 package workload
 
+import (
+	"encoding/json"
+	"errors"
+)
+
 type WorkloadType string
 
 const (
 	WorkloadTypeJob     = "Job"
 	WorkloadTypeService = "Service"
 )
 
+type WorkloadHeader struct {
+	Type WorkloadType `yaml:"_type" json:"_type"`
+}
+
 type Workload struct {
-	Type     WorkloadType `yaml:"type" json:"type"`
-	*Service `yaml:",inline" json:",inline"`
-	*Job     `yaml:",inline" json:",inline"`
+	WorkloadHeader `yaml:",inline" json:",inline"`
+	*Service       `yaml:",inline" json:",inline"`
+	*Job           `yaml:",inline" json:",inline"`
+}
+
+func (w Workload) MarshalJSON() ([]byte, error) {
+	switch w.Type {
+	case WorkloadTypeService:
+		return json.Marshal(struct {
+			WorkloadHeader `yaml:",inline" json:",inline"`
+			*Service       `json:",inline"`
+		}{
+			WorkloadHeader: WorkloadHeader{w.Type},
+			Service:        w.Service,
+		})
+	case WorkloadTypeJob:
+		return json.Marshal(struct {
+			WorkloadHeader `yaml:",inline" json:",inline"`
+			*Job           `json:",inline"`
+		}{
+			WorkloadHeader: WorkloadHeader{w.Type},
+			Job:            w.Job,
+		})
+	default:
+		return nil, errors.New("unknown workload type")
+	}
+}
+
+func (w *Workload) UnmarshalJSON(data []byte) error {
+	var workloadData WorkloadHeader
+	err := json.Unmarshal(data, &workloadData)
+	if err != nil {
+		return err
+	}
+
+	w.Type = workloadData.Type
+	switch w.Type {
+	case WorkloadTypeJob:
+		var v Job
+		err = json.Unmarshal(data, &v)
+		w.Job = &v
+	case WorkloadTypeService:
+		var v Service
+		err = json.Unmarshal(data, &v)
+		w.Service = &v
+	default:
+		err = errors.New("unknown workload type")
+	}
+
+	return err
 }

pkg/models/appconfiguration/workload/workload_test.go

@@ -0,0 +1,87 @@
+package workload
+
+import (
+	"encoding/json"
+	"errors"
+	"testing"
+)
+
+func TestWorkload_MarshalJSON(t *testing.T) {
+	workload := Workload{
+		WorkloadHeader: WorkloadHeader{
+			Type: WorkloadTypeService,
+		},
+		Service: &Service{
+			WorkloadBase: WorkloadBase{
+				Replicas: 2,
+				Labels: map[string]string{
+					"app": "my-service",
+				},
+			},
+		},
+		Job: &Job{
+			Schedule: "* * * * *",
+		},
+	}
+
+	expected := `{"_type":"Service","replicas":2,"labels":{"app":"my-service"}}`
+
+	data, err := json.Marshal(workload)
+	if err != nil {
+		t.Errorf("Error while marshaling workload: %v", err)
+	}
+
+	if string(data) != expected {
+		t.Errorf("Expected marshaled JSON: %s, got: %s", expected, string(data))
+	}
+}
+
+func TestWorkload_UnmarshalJSON(t *testing.T) {
+	jsonData := `{"_type":"Service","replicas":1,"labels":{},"annotations":{},"dirs":{},"schedule":"* * * * *"}`
+	expected := Workload{
+		WorkloadHeader: WorkloadHeader{
+			Type: WorkloadTypeService,
+		},
+		Service: &Service{
+			WorkloadBase: WorkloadBase{
+				Replicas:    1,
+				Labels:      map[string]string{},
+				Annotations: map[string]string{},
+				Dirs:        map[string]string{},
+			},
+		},
+	}
+
+	var actual Workload
+	err := json.Unmarshal([]byte(jsonData), &actual)
+	if err != nil {
+		t.Errorf("Error while unmarshaling JSON: %v", err)
+	}
+
+	if actual.Type != expected.Type {
+		t.Errorf("Expected workload type: %s, got: %s", expected.Type, actual.Type)
+	}
+
+	if actual.Service == nil {
+		t.Errorf("Expected service is not nil, got: %v", expected.Service)
+	}
+
+	if actual.Job != nil {
+		t.Errorf("Expected job is nil, got: %v", expected.Job)
+	}
+}
+
+func TestWorkload_UnmarshalJSON_UnknownType(t *testing.T) {
+	jsonData := `{"_type":"Unknown","replicas":1,"labels":{},"annotations":{},"dirs":{},"schedule":"* * * * *"}`
+
+	var workload Workload
+	err := json.Unmarshal([]byte(jsonData), &workload)
+	if err == nil {
+		t.Error("Expected error for unknown workload type")
+	}
+
+	expectedError := errors.New("unknown workload type")
+	if err.Error() != expectedError.Error() {
+		t.Errorf("Expected error: %v, got: %v", expectedError, err)
+	}
+}