From 051809f7deed8702dcb22002bbab1f7baff6e458 Mon Sep 17 00:00:00 2001
From: "bin.hu" <bin.hu@ringcentral.com>
Date: Fri, 15 Sep 2023 15:53:05 +0800
Subject: [PATCH] nginx canary weight support headers

Signed-off-by: bin.hu <bin.hu@ringcentral.com>
---
 pkg/router/ingress.go      |  9 +++---
 pkg/router/ingress_test.go | 59 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/pkg/router/ingress.go b/pkg/router/ingress.go
index 7fd9e4207..431546355 100644
--- a/pkg/router/ingress.go
+++ b/pkg/router/ingress.go
@@ -135,7 +135,7 @@ func (i *IngressRouter) GetRoutes(canary *flaggerv1.Canary) (
 	}
 
 	// A/B testing
-	if len(canary.GetAnalysis().Match) > 0 {
+	if len(canary.GetAnalysis().Match) > 0 && canary.GetAnalysis().Iterations > 0 {
 		for k := range canaryIngress.Annotations {
 			if k == i.GetAnnotationWithPrefix("canary-by-cookie") || k == i.GetAnnotationWithPrefix("canary-by-header") {
 				return 0, 100, false, nil
@@ -156,7 +156,6 @@ func (i *IngressRouter) GetRoutes(canary *flaggerv1.Canary) (
 			break
 		}
 	}
-
 	primaryWeight = 100 - canaryWeight
 	mirrored = false
 	return
@@ -176,7 +175,8 @@ func (i *IngressRouter) SetRoutes(
 
 	iClone := canaryIngress.DeepCopy()
 
-	// A/B testing
+	// if iterations > 0 ,it's A/B testing.
+	// else it's a weight canary with match header
 	if len(canary.GetAnalysis().Match) > 0 {
 		var cookie, header, headerValue, headerRegex string
 		for _, m := range canary.GetAnalysis().Match {
@@ -192,7 +192,8 @@ func (i *IngressRouter) SetRoutes(
 		}
 
 		iClone.Annotations = i.makeHeaderAnnotations(iClone.Annotations, header, headerValue, headerRegex, cookie)
-	} else {
+	}
+	if canary.GetAnalysis().Iterations == 0 {
 		// canary
 		iClone.Annotations[i.GetAnnotationWithPrefix("canary-weight")] = fmt.Sprintf("%v", canaryWeight)
 	}
diff --git a/pkg/router/ingress_test.go b/pkg/router/ingress_test.go
index f330777c2..9c9c19c4e 100644
--- a/pkg/router/ingress_test.go
+++ b/pkg/router/ingress_test.go
@@ -100,6 +100,65 @@ func TestIngressRouter_GetSetRoutes(t *testing.T) {
 	assert.Equal(t, "0", inCanary.Annotations[canaryWeightAn])
 }
 
+func TestIngressRouter_GetSetRoutes_WithMatch(t *testing.T) {
+	mocks := newFixture(nil)
+	router := &IngressRouter{
+		logger:            mocks.logger,
+		kubeClient:        mocks.kubeClient,
+		annotationsPrefix: "prefix1.nginx.ingress.kubernetes.io",
+	}
+	mocks.ingressCanary.Spec.Analysis.Match = []istiov1alpha3.HTTPMatchRequest{
+		{
+			Headers: map[string]istiov1alpha1.StringMatch{
+				"x-canary": {
+					Exact: "always",
+				},
+			},
+		},
+	}
+	err := router.Reconcile(mocks.ingressCanary)
+	require.NoError(t, err)
+
+	p, c, m, err := router.GetRoutes(mocks.ingressCanary)
+	require.NoError(t, err)
+
+	p = 50
+	c = 50
+	m = false
+
+	err = router.SetRoutes(mocks.ingressCanary, p, c, m)
+	require.NoError(t, err)
+
+	canaryAn := "prefix1.nginx.ingress.kubernetes.io/canary"
+	canaryWeightAn := "prefix1.nginx.ingress.kubernetes.io/canary-weight"
+	canaryByHeaderAn := "prefix1.nginx.ingress.kubernetes.io/canary-by-header"
+	canaryByHeaderValueAn := "prefix1.nginx.ingress.kubernetes.io/canary-by-header-value"
+
+	canaryName := fmt.Sprintf("%s-canary", mocks.ingressCanary.Spec.IngressRef.Name)
+	inCanary, err := router.kubeClient.NetworkingV1().Ingresses("default").Get(context.TODO(), canaryName, metav1.GetOptions{})
+	require.NoError(t, err)
+
+	// test rollout
+	assert.Equal(t, "true", inCanary.Annotations[canaryAn])
+	assert.Equal(t, "50", inCanary.Annotations[canaryWeightAn])
+	assert.Equal(t, "x-canary", inCanary.Annotations[canaryByHeaderAn])
+	assert.Equal(t, "always", inCanary.Annotations[canaryByHeaderValueAn])
+
+	p = 100
+	c = 0
+	m = false
+
+	err = router.SetRoutes(mocks.ingressCanary, p, c, m)
+	require.NoError(t, err)
+
+	inCanary, err = router.kubeClient.NetworkingV1().Ingresses("default").Get(context.TODO(), canaryName, metav1.GetOptions{})
+	require.NoError(t, err)
+
+	// test promotion
+	assert.Equal(t, "true", inCanary.Annotations[canaryAn])
+	assert.Equal(t, "0", inCanary.Annotations[canaryWeightAn])
+}
+
 func TestIngressRouter_ABTest(t *testing.T) {
 	mocks := newFixture(nil)
 	router := &IngressRouter{