From 61360300640935056362e75988fad1c2f976a275 Mon Sep 17 00:00:00 2001
From: njucjc <njucjc@gmail.com>
Date: Mon, 8 Nov 2021 15:09:25 +0800
Subject: [PATCH] add node affinity

---
 tests/config_test.go     |  2 +-
 tests/connection_test.go | 19 +++++++-------
 tests/prepare_test.go    | 57 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 11 deletions(-)

diff --git a/tests/config_test.go b/tests/config_test.go
index 20fb753d..8b85e76f 100644
--- a/tests/config_test.go
+++ b/tests/config_test.go
@@ -15,7 +15,7 @@ import (
 
 var enableTrunk bool
 var enablePolicy bool
-var testNamespace = "network-test" + strconv.FormatInt(time.Now().Unix(), 10)
+var testNamespace = "network-test-" + strconv.FormatInt(time.Now().Unix(), 10)
 
 func init() {
 	flag.BoolVar(&enableTrunk, "trunk", false, "install trunk policy")
diff --git a/tests/connection_test.go b/tests/connection_test.go
index 5e78efd6..5f957252 100644
--- a/tests/connection_test.go
+++ b/tests/connection_test.go
@@ -277,9 +277,10 @@ func (s *ConnectionTestSuite) SetupSuite() {
 func (s *ConnectionTestSuite) TearDownSuite() {
 
 	if s.err != nil {
-		s.T().Errorf("skip tear down resource in namespace %s, because of an error occurred.", testNamespace)
+		s.T().Error(errors.Wrapf(s.err, "skip tear down resource in namespace %s, because of an error occurred.", testNamespace))
 		return
 	}
+
 	ctx := context.Background()
 	if enablePolicy {
 		s.T().Logf("delete %s", networkPolicy.Name)
@@ -374,9 +375,8 @@ func (s *ConnectionTestSuite) TestPod2Pod() {
 				for _, ip := range podIPs(&dst) {
 					addr := net.JoinHostPort(ip, "80")
 					l := fmt.Sprintf("src %s -> dst %s", podInfo(&src), addr)
-					var stdErrOut []byte
-					_, stdErrOut, s.err = s.ExecHTTPGet(src.Namespace, src.Name, curlAddr(addr))
-					s.Expected(c.Status, stdErrOut, s.err, l)
+					_, stdErrOut, err := s.ExecHTTPGet(src.Namespace, src.Name, curlAddr(addr))
+					s.Expected(c.Status, stdErrOut, err, l)
 				}
 			}
 		}
@@ -443,9 +443,8 @@ func (s *ConnectionTestSuite) TestPod2ServiceIP() {
 
 				for _, addr := range addrs {
 					l := fmt.Sprintf("src %s -> dst svc name %s, addr %s", podInfo(&src), svc.Name, addr)
-					var stdErrOut []byte
-					_, stdErrOut, s.err = s.ExecHTTPGet(src.Namespace, src.Name, curlAddr(addr))
-					s.Expected(c.Status, stdErrOut, s.err, l)
+					_, stdErrOut, err := s.ExecHTTPGet(src.Namespace, src.Name, curlAddr(addr))
+					s.Expected(c.Status, stdErrOut, err, l)
 				}
 			}
 		}
@@ -473,9 +472,8 @@ func (s *ConnectionTestSuite) TestPod2ServiceName() {
 		for _, src := range srcPods {
 			for _, svc := range dstServices {
 				l := fmt.Sprintf("src %s -> dst svc name %s", podInfo(&src), svc.Name)
-				var stdErrOut []byte
-				_, stdErrOut, s.err = s.ExecHTTPGet(src.Namespace, src.Name, svc.Name)
-				s.Expected(c.Status, stdErrOut, s.err, l)
+				_, stdErrOut, err := s.ExecHTTPGet(src.Namespace, src.Name, svc.Name)
+				s.Expected(c.Status, stdErrOut, err, l)
 			}
 		}
 	}
@@ -486,6 +484,7 @@ func (s *ConnectionTestSuite) Expected(status bool, stdErrOut []byte, err error,
 		if assert.NoError(s.T(), err, msg) && assert.Equal(s.T(), 0, len(stdErrOut), msg) {
 			s.T().Logf(msg + ", test pass")
 		} else {
+			s.err = err
 			s.T().Error(errors.Wrapf(err, "%s, test failed, expected connection success, but connection failure", msg))
 		}
 	} else {
diff --git a/tests/prepare_test.go b/tests/prepare_test.go
index b294dd68..6428a6a2 100644
--- a/tests/prepare_test.go
+++ b/tests/prepare_test.go
@@ -149,6 +149,25 @@ func EnsureDaemonSet(ctx context.Context, cs kubernetes.Interface, cfg PodResCon
 							ImagePullPolicy: corev1.PullAlways,
 						},
 					},
+					Affinity: &corev1.Affinity{
+						NodeAffinity: &corev1.NodeAffinity{
+							RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+								NodeSelectorTerms: []corev1.NodeSelectorTerm{
+									{
+										MatchExpressions: []corev1.NodeSelectorRequirement{
+											{
+												Key:      "type",
+												Operator: corev1.NodeSelectorOpNotIn,
+												Values: []string{
+													"virtual-kubelet",
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
 				},
 			},
 		},
@@ -204,6 +223,25 @@ func EnsureDeployment(ctx context.Context, cs kubernetes.Interface, cfg PodResCo
 							ImagePullPolicy: corev1.PullAlways,
 						},
 					},
+					Affinity: &corev1.Affinity{
+						NodeAffinity: &corev1.NodeAffinity{
+							RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+								NodeSelectorTerms: []corev1.NodeSelectorTerm{
+									{
+										MatchExpressions: []corev1.NodeSelectorRequirement{
+											{
+												Key:      "type",
+												Operator: corev1.NodeSelectorOpNotIn,
+												Values: []string{
+													"virtual-kubelet",
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
 				},
 			},
 		},
@@ -256,6 +294,25 @@ func EnsureStatefulSet(ctx context.Context, cs kubernetes.Interface, cfg PodResC
 							ImagePullPolicy: corev1.PullAlways,
 						},
 					},
+					Affinity: &corev1.Affinity{
+						NodeAffinity: &corev1.NodeAffinity{
+							RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+								NodeSelectorTerms: []corev1.NodeSelectorTerm{
+									{
+										MatchExpressions: []corev1.NodeSelectorRequirement{
+											{
+												Key:      "type",
+												Operator: corev1.NodeSelectorOpNotIn,
+												Values: []string{
+													"virtual-kubelet",
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
 				},
 			},
 		},