Fix: loading definitions if lack in the revision when calculate dispatch stage of trait. (#5644)

* small fix for dispatch stage Signed-off-by: 楚岳 <wangyike.wyk@alibaba-inc.com> revert useless commit Signed-off-by: 楚岳 <wangyike.wyk@alibaba-inc.com> small fix Signed-off-by: 楚岳 <wangyike.wyk@alibaba-inc.com> * add test case Signed-off-by: 楚岳 <wangyike.wyk@alibaba-inc.com> --------- Signed-off-by: 楚岳 <wangyike.wyk@alibaba-inc.com>
2026-08-18 20:17:04 +00:00 · 2023-03-13 11:22:44 +08:00
parent 4edb21782f
commit a36eb4f2e2
2 changed files with 103 additions and 12 deletions
@@ -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))
+	})
+})
@@ -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
+}