Address comments

Signed-off-by: graysonwu <[email protected]>
antrea-io · Jan 2, 2024 · b427b77 · b427b77
1 parent 588385b
commit b427b77
Showing 1 changed file with 28 additions and 28 deletions.
diff --git a/pkg/agent/controller/networkpolicy/reject_test.go b/pkg/agent/controller/networkpolicy/reject_test.go
@@ -473,7 +473,7 @@ func Test_getRejectPacketOutMutateFunc(t *testing.T) {
 		{
 			name: "RejectLocalToRemoteNoFlexibleIPAM,ExternalNode",
 			args: args{
-				rejectType:        RejectPodLocalToRemote,
+				rejectType:        rejectPodLocalToRemote,
 				nodeType:          config.ExternalNode,
 				isFlexibleIPAMSrc: false,
 				isFlexibleIPAMDst: false,
@@ -559,20 +559,31 @@ func Test_handleRejectRequest(t *testing.T) {
 		name              string
 		packetIn          *ofctrl.PacketIn
 		antreaProxyEnable bool
-		expectFunc        func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore)
+		prepareFunc       func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore)
 	}{
 		{
-			name:              "IPv4,TCP,antreaProxy,podLocal",
+			name:              "RejectPodLocalIPv4TCPTrafficWithAntreaProxy",
 			packetIn:          genPacketIn(fakeSrcMac, fakeDstMac, fakeSrcIPv4, fakeDstIPv4, fakeSrcPort, fakeDstPort, fakeTCPSeqNum, false, true, []openflow15.MatchField{*openflow15.NewInPortField(uint32(fakeOFPort1))}),
 			antreaProxyEnable: true,
-			expectFunc: func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore) {
+			prepareFunc: func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore) {
+				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(fakeSIface, true)
+				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(fakeDIface, true)
+				mockClient.EXPECT().SendTCPPacketOut(fakeDstMac, fakeSrcMac, fakeDstIPv4, fakeSrcIPv4, gomock.Any(), gomock.Any(), false, fakeDstPort, fakeSrcPort, gomock.Any(), fakeTCPSeqNum+1, gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil)
+			},
+		},
+		{
+			name:              "RejectPodLocalIPv4TCPTrafficNoAntreaProxy",
+			packetIn:          genPacketIn(fakeSrcMac, fakeDstMac, fakeSrcIPv4, fakeDstIPv4, fakeSrcPort, fakeDstPort, fakeTCPSeqNum, false, true, nil),
+			antreaProxyEnable: false,
+			prepareFunc: func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore) {
 				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(fakeSIface, true)
 				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(fakeDIface, true)
+				mockIStore.EXPECT().GetInterfacesByType(gomock.Any()).Return([]*interfacestore.InterfaceConfig{fakeGWIface})
 				mockClient.EXPECT().SendTCPPacketOut(fakeDstMac, fakeSrcMac, fakeDstIPv4, fakeSrcIPv4, gomock.Any(), gomock.Any(), false, fakeDstPort, fakeSrcPort, gomock.Any(), fakeTCPSeqNum+1, gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil)
 			},
 		},
 		{
-			name: "IPv4,TCP,antreaProxy,serviceRemoteToExternal",
+			name: "RejectServiceRemoteToExternalIPv4TCPTrafficWithAntreaProxy",
 			packetIn: genPacketIn(fakeSrcMac, fakeDstMac, fakeSrcIPv4, fakeDstIPv4, fakeSrcPort, fakeDstPort, fakeTCPSeqNum, false, true, []openflow15.MatchField{{
 				// EpSelectedRegMark
 				Class:          openflow15.OXM_CLASS_PACKET_REGS,
@@ -587,54 +598,43 @@ func Test_handleRejectRequest(t *testing.T) {
 				Mask: nil,
 			}}),
 			antreaProxyEnable: true,
-			expectFunc: func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore) {
+			prepareFunc: func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore) {
 				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(nil, false)
 				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(nil, false)
 				mockClient.EXPECT().SendTCPPacketOut(fakeDstMac, openflow.GlobalVirtualMAC.String(), fakeDstIPv4, fakeSrcIPv4, gomock.Any(), gomock.Any(), false, fakeDstPort, fakeSrcPort, gomock.Any(), fakeTCPSeqNum+1, gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil)
 			},
 		},
 		{
-			name:              "IPv4,UDP,antreaProxy,podLocal",
+			name:              "RejectPodLocalIPv4UDPTrafficWithAntreaProxy",
 			packetIn:          genPacketIn(fakeSrcMac, fakeDstMac, fakeSrcIPv4, fakeDstIPv4, fakeSrcPort, fakeDstPort, 0, false, false, []openflow15.MatchField{*openflow15.NewInPortField(uint32(fakeOFPort1))}),
 			antreaProxyEnable: true,
-			expectFunc: func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore) {
+			prepareFunc: func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore) {
 				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(fakeSIface, true)
 				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(fakeDIface, true)
 				mockClient.EXPECT().SendICMPPacketOut(fakeDstMac, fakeSrcMac, fakeDstIPv4, fakeSrcIPv4, gomock.Any(), gomock.Any(), false, uint8(3), uint8(10), gomock.Any(), gomock.Any()).Times(1).Return(nil)
 			},
 		},
 		{
-			name:              "IPv4,TCP,antreaProxyDisable,podLocal",
-			packetIn:          genPacketIn(fakeSrcMac, fakeDstMac, fakeSrcIPv4, fakeDstIPv4, fakeSrcPort, fakeDstPort, fakeTCPSeqNum, false, true, nil),
-			antreaProxyEnable: false,
-			expectFunc: func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore) {
-				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(fakeSIface, true)
-				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(fakeDIface, true)
-				mockIStore.EXPECT().GetInterfacesByType(gomock.Any()).Return([]*interfacestore.InterfaceConfig{fakeGWIface})
-				mockClient.EXPECT().SendTCPPacketOut(fakeDstMac, fakeSrcMac, fakeDstIPv4, fakeSrcIPv4, gomock.Any(), gomock.Any(), false, fakeDstPort, fakeSrcPort, gomock.Any(), fakeTCPSeqNum+1, gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil)
-			},
-		},
-		{
-			name:              "IPv6,TCP,antreaProxy,podLocal",
+			name:              "RejectPodLocalIPv6TCPTrafficWithAntreaProxy",
 			packetIn:          genPacketIn(fakeSrcMac, fakeDstMac, fakeSrcIPv6, fakeDstIPv6, fakeSrcPort, fakeDstPort, fakeTCPSeqNum, true, true, nil),
 			antreaProxyEnable: true,
-			expectFunc: func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore) {
+			prepareFunc: func(mockClient *oftesting.MockClient, mockIStore *storetesting.MockInterfaceStore) {
 				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(fakeSIface, true)
 				mockIStore.EXPECT().GetInterfaceByIP(gomock.Any()).Return(fakeDIface, true)
 				mockClient.EXPECT().SendTCPPacketOut(fakeDstMac, fakeSrcMac, fakeDstIPv6, fakeSrcIPv6, gomock.Any(), gomock.Any(), true, fakeDstPort, fakeSrcPort, gomock.Any(), fakeTCPSeqNum+1, gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(1).Return(nil)
 			},
 		},
 	}
-	ctrl := gomock.NewController(t)
-	controller, _, _ := newTestController()
-	mockClient := oftesting.NewMockClient(ctrl)
-	mockIStore := storetesting.NewMockInterfaceStore(ctrl)
-	controller.ofClient = mockClient
-	controller.ifaceStore = mockIStore
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			controller, _, _ := newTestController()
+			mockClient := oftesting.NewMockClient(ctrl)
+			mockIStore := storetesting.NewMockInterfaceStore(ctrl)
+			controller.ofClient = mockClient
+			controller.ifaceStore = mockIStore
 			controller.antreaProxyEnabled = tt.antreaProxyEnable
-			tt.expectFunc(mockClient, mockIStore)
+			tt.prepareFunc(mockClient, mockIStore)
 			assert.NoError(t, controller.rejectRequest(tt.packetIn))
 		})
 	}