diff --git a/pkg/controller/core.oam.dev/v1alpha2/application/dispatch_test.go b/pkg/controller/core.oam.dev/v1alpha2/application/dispatch_test.go new file mode 100644 index 000000000..cf6f8a8b1 --- /dev/null +++ b/pkg/controller/core.oam.dev/v1alpha2/application/dispatch_test.go @@ -0,0 +1,75 @@ +/* +Copyright 2022 The KubeVela Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package application + +import ( + "context" + + . "github.com/onsi/ginkgo" + . "github.com/onsi/gomega" + v1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + "github.com/oam-dev/kubevela/apis/core.oam.dev/v1beta1" + velatypes "github.com/oam-dev/kubevela/apis/types" +) + +var _ = Describe("Test dispatch stage", func() { + BeforeEach(func() { + traitDefinition := v1beta1.TraitDefinition{ + ObjectMeta: v1.ObjectMeta{ + Name: "kruise-rollout", + Namespace: velatypes.DefaultKubeVelaNS, + }, + Spec: v1beta1.TraitDefinitionSpec{ + Stage: v1beta1.PreDispatch, + }, + } + Expect(k8sClient.Create(context.Background(), &traitDefinition)).Should(BeNil()) + }) + + It("Test get dispatch stage from trait", func() { + appRev := v1beta1.ApplicationRevision{ + Spec: v1beta1.ApplicationRevisionSpec{ + ApplicationRevisionCompressibleFields: v1beta1.ApplicationRevisionCompressibleFields{ + TraitDefinitions: map[string]v1beta1.TraitDefinition{ + "gateway": { + Spec: v1beta1.TraitDefinitionSpec{ + Stage: v1beta1.PostDispatch, + }, + }, + "hpa": { + Spec: v1beta1.TraitDefinitionSpec{}, + }, + }, + }, + }, + } + + stage, err := getTraitDispatchStage(k8sClient, "kruise-rollout", &appRev) + Expect(err).Should(BeNil()) + Expect(stage).Should(BeEquivalentTo(PreDispatch)) + stage, err = getTraitDispatchStage(k8sClient, "gateway", &appRev) + Expect(err).Should(BeNil()) + Expect(stage).Should(BeEquivalentTo(PostDispatch)) + stage, err = getTraitDispatchStage(k8sClient, "hpa", &appRev) + Expect(err).Should(BeNil()) + Expect(stage).Should(BeEquivalentTo(DefaultDispatch)) + stage, err = getTraitDispatchStage(k8sClient, "not-exist", &appRev) + Expect(err).ShouldNot(BeNil()) + Expect(stage).Should(BeEquivalentTo(DefaultDispatch)) + }) +}) diff --git a/pkg/controller/core.oam.dev/v1alpha2/application/dispatcher.go b/pkg/controller/core.oam.dev/v1alpha2/application/dispatcher.go index 3b81559ea..666a7e7df 100644 --- a/pkg/controller/core.oam.dev/v1alpha2/application/dispatcher.go +++ b/pkg/controller/core.oam.dev/v1alpha2/application/dispatcher.go @@ -21,12 +21,14 @@ import ( "sort" "strings" + "github.com/pkg/errors" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "sigs.k8s.io/controller-runtime/pkg/client" + "github.com/oam-dev/kubevela/apis/core.oam.dev/common" "github.com/oam-dev/kubevela/pkg/cue/definition" "github.com/oam-dev/kubevela/pkg/oam" - - "github.com/pkg/errors" - "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + oamutil "github.com/oam-dev/kubevela/pkg/oam/util" "github.com/oam-dev/kubevela/apis/core.oam.dev/v1beta1" "github.com/oam-dev/kubevela/pkg/appfile" @@ -177,15 +179,9 @@ func (h *AppHandler) generateDispatcher(appRev *v1beta1.ApplicationRevision, rea traitType = splitName } } - if trait, ok := appRev.Spec.TraitDefinitions[traitType]; ok { - _stageType := trait.Spec.Stage - if len(_stageType) == 0 { - _stageType = v1beta1.DefaultDispatch - } - stageType, err = ParseStageType(string(_stageType)) - if err != nil { - return nil, err - } + stageType, err = getTraitDispatchStage(h.r.Client, traitType, appRev) + if err != nil { + return nil, err } } traitStageMap[stageType] = append(traitStageMap[stageType], readyTrait) @@ -213,3 +209,23 @@ func (h *AppHandler) generateDispatcher(appRev *v1beta1.ApplicationRevision, rea } return manifestDispatchers, nil } + +func getTraitDispatchStage(client client.Client, traitType string, appRev *v1beta1.ApplicationRevision) (StageType, error) { + trait, ok := appRev.Spec.TraitDefinitions[traitType] + if !ok { + trait = v1beta1.TraitDefinition{} + err := oamutil.GetCapabilityDefinition(context.Background(), client, &trait, traitType) + if err != nil { + return DefaultDispatch, err + } + } + _stageType := trait.Spec.Stage + if len(_stageType) == 0 { + _stageType = v1beta1.DefaultDispatch + } + stageType, err := ParseStageType(string(_stageType)) + if err != nil { + return DefaultDispatch, err + } + return stageType, nil +}