diff --git a/pkg/models/appconfiguration/workload/common.go b/pkg/models/appconfiguration/workload/common.go index 987513c1..937c5c1c 100644 --- a/pkg/models/appconfiguration/workload/common.go +++ b/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. diff --git a/pkg/models/appconfiguration/workload/job.go b/pkg/models/appconfiguration/workload/job.go index 8fa5c47a..ad767e18 100644 --- a/pkg/models/appconfiguration/workload/job.go +++ b/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"` } diff --git a/pkg/models/appconfiguration/workload/service.go b/pkg/models/appconfiguration/workload/service.go index a60c3e23..56b0ec7f 100644 --- a/pkg/models/appconfiguration/workload/service.go +++ b/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"` } diff --git a/pkg/models/appconfiguration/workload/workload.go b/pkg/models/appconfiguration/workload/workload.go index fa007404..0c2d3ac7 100644 --- a/pkg/models/appconfiguration/workload/workload.go +++ b/pkg/models/appconfiguration/workload/workload.go @@ -1,5 +1,10 @@ package workload +import ( + "encoding/json" + "errors" +) + type WorkloadType string const ( @@ -7,8 +12,59 @@ const ( 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 } diff --git a/pkg/models/appconfiguration/workload/workload_test.go b/pkg/models/appconfiguration/workload/workload_test.go new file mode 100644 index 00000000..462a4aac --- /dev/null +++ b/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) + } +}