Add unit tests for Flow Aggregator API handlers (antrea-io#4179)

Signed-off-by: Yanjun Zhou <zhouya@vmware.com>

hack/update-codegen-dockerized.sh

@@ -62,6 +62,7 @@ function generate_mocks {
     "pkg/ovs/ovsconfig OVSBridgeClient testing"
     "pkg/ovs/ovsctl OVSCtlClient testing"
     "pkg/querier AgentNetworkPolicyInfoQuerier,AgentMulticastInfoQuerier testing"
+    "pkg/flowaggregator/querier FlowAggregatorQuerier testing"
     "third_party/proxy Provider testing"
   )
 

pkg/flowaggregator/apiserver/handlers/flowrecords/handler_test.go

@@ -0,0 +1,176 @@
+// Copyright 2022 Antrea 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 flowrecords
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+	ipfixintermediate "github.com/vmware/go-ipfix/pkg/intermediate"
+
+	queriertest "antrea.io/antrea/pkg/flowaggregator/querier/testing"
+)
+
+var (
+	record1 = map[string]interface{}{
+		"sourceIPv4Address":          "10.0.0.1",
+		"destinationIPv4Address":     "10.0.0.2",
+		"sourceTransportPort":        float64(8080),
+		"destinationTransportPort":   float64(3700),
+		"protocolIdentifier":         float64(6),
+		"sourcePodName":              "test-pod-a",
+		"destinationPodName":         "test-pod-b",
+		"sourcePodNamespace":         "test-namespace-a",
+		"destinationPodNamespace":    "test-namespace-b",
+		"destinationServicePortName": "",
+	}
+	recordTableRows1 = []string{
+		"10.0.0.1", "10.0.0.2", "8080", "3700", "6", "test-pod-a", "test-pod-b", "test-namespace-a", "test-namespace-b", "",
+	}
+	record2 = map[string]interface{}{
+		"sourceIPv4Address":          "10.0.0.2",
+		"destinationIPv4Address":     "10.0.0.1",
+		"sourceTransportPort":        float64(3701),
+		"destinationTransportPort":   float64(8080),
+		"protocolIdentifier":         float64(6),
+		"sourcePodName":              "test-pod-b",
+		"destinationPodName":         "test-pod-a",
+		"sourcePodNamespace":         "test-namespace-b",
+		"destinationPodNamespace":    "test-namespace-a",
+		"destinationServicePortName": "",
+	}
+	recordTableRows2 = []string{
+		"10.0.0.2", "10.0.0.1", "3701", "8080", "6", "test-pod-b", "test-pod-a", "test-namespace-b", "test-namespace-a", "",
+	}
+	record3 = map[string]interface{}{
+		"sourceIPv6Address":          "2001:ce18:4d:1c5::3ac:2d1",
+		"destinationIPv6Address":     "2001:ce18:4d:1c5::3ac:2d2",
+		"sourceTransportPort":        float64(8080),
+		"destinationTransportPort":   float64(3700),
+		"protocolIdentifier":         float64(6),
+		"sourcePodName":              "test-pod-c",
+		"destinationPodName":         "test-pod-d",
+		"sourcePodNamespace":         "test-namespace-c",
+		"destinationPodNamespace":    "test-namespace-d",
+		"destinationServicePortName": "",
+	}
+	recordTableRows3 = []string{
+		"2001:ce18:4d:1c5::3ac:2d1", "2001:ce18:4d:1c5::3ac:2d2", "8080", "3700", "6", "test-pod-c", "test-pod-d", "test-namespace-c", "test-namespace-d", "",
+	}
+)
+
+type testCase struct {
+	name              string
+	records           []map[string]interface{}
+	query             string
+	flowKey           *ipfixintermediate.FlowKey
+	expectedStatus    int
+	expectedResponse  []Response
+	expectedTableRows [][]string
+}
+
+func TestGetFlowRecordsQuery(t *testing.T) {
+	testCases := []testCase{
+		{
+			name:              "Get all records",
+			records:           []map[string]interface{}{record1, record2, record3},
+			query:             "",
+			flowKey:           nil,
+			expectedStatus:    http.StatusOK,
+			expectedResponse:  []Response{record1, record2, record3},
+			expectedTableRows: [][]string{recordTableRows1, recordTableRows2, recordTableRows3},
+		},
+		{
+			name:    "Get records by IP address",
+			records: []map[string]interface{}{record2},
+			query:   "?srcip=10.0.0.2&&dstip=10.0.0.1",
+			flowKey: &ipfixintermediate.FlowKey{
+				SourceAddress:      "10.0.0.2",
+				DestinationAddress: "10.0.0.1",
+			},
+			expectedStatus:    http.StatusOK,
+			expectedResponse:  []Response{record2},
+			expectedTableRows: [][]string{recordTableRows2},
+		},
+		{
+			name:    "Get records by ports",
+			records: []map[string]interface{}{record1, record3},
+			query:   "?srcport=8080",
+			flowKey: &ipfixintermediate.FlowKey{
+				SourcePort: 8080,
+			},
+			expectedStatus:    http.StatusOK,
+			expectedResponse:  []Response{record1, record3},
+			expectedTableRows: [][]string{recordTableRows1, recordTableRows3},
+		},
+		{
+			name:  "Records not found",
+			query: "?srcip=10.0.0.10",
+			flowKey: &ipfixintermediate.FlowKey{
+				SourceAddress: "10.0.0.10",
+			},
+			expectedStatus: http.StatusOK,
+		},
+		{
+			name:           "Illegal protocol",
+			query:          "?proto=tcp",
+			expectedStatus: http.StatusNotFound,
+		},
+		{
+			name:           "Illegal source port",
+			query:          "?srcport=tcp-port",
+			expectedStatus: http.StatusNotFound,
+		},
+		{
+			name:           "Illegal destination port",
+			query:          "?dstport=tcp-port",
+			expectedStatus: http.StatusNotFound,
+		},
+	}
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			faq := queriertest.NewMockFlowAggregatorQuerier(ctrl)
+			faq.EXPECT().GetFlowRecords(tc.flowKey).Return(tc.records).AnyTimes()
+
+			handler := HandleFunc(faq)
+			req, err := http.NewRequest(http.MethodGet, tc.query, nil)
+			assert.Nil(t, err)
+			recorder := httptest.NewRecorder()
+			handler.ServeHTTP(recorder, req)
+			assert.Equal(t, tc.expectedStatus, recorder.Code)
+
+			if tc.expectedStatus == http.StatusOK {
+				var received []Response
+				err = json.Unmarshal(recorder.Body.Bytes(), &received)
+				assert.Nil(t, err)
+				assert.ElementsMatch(t, tc.expectedResponse, received)
+				var receivedTableRows [][]string
+				for _, r := range received {
+					receivedTableRows = append(receivedTableRows, r.GetTableRow(0))
+				}
+				assert.ElementsMatch(t, tc.expectedTableRows, receivedTableRows)
+			}
+		})
+
+	}
+
+}

pkg/flowaggregator/apiserver/handlers/recordmetrics/handler.go

@@ -45,7 +45,7 @@ func HandleFunc(faq querier.FlowAggregatorQuerier) http.HandlerFunc {
 		err := json.NewEncoder(w).Encode(metricsResponse)
 		if err != nil {
 			w.WriteHeader(http.StatusInternalServerError)
-			klog.Errorf("Error when encoding AntreaAgentInfo to json: %v", err)
+			klog.Errorf("Error when encoding record metrics to json: %v", err)
 		}
 	}
 }

pkg/flowaggregator/apiserver/handlers/recordmetrics/handler_test.go

@@ -0,0 +1,60 @@
+// Copyright 2022 Antrea 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 recordmetrics
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+
+	"antrea.io/antrea/pkg/flowaggregator/querier"
+	queriertest "antrea.io/antrea/pkg/flowaggregator/querier/testing"
+)
+
+func TestRecordMetricsQuery(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	faq := queriertest.NewMockFlowAggregatorQuerier(ctrl)
+	faq.EXPECT().GetRecordMetrics().Return(querier.Metrics{
+		NumRecordsExported: 20,
+		NumRecordsReceived: 15,
+		NumFlows:           30,
+		NumConnToCollector: 1,
+	})
+
+	handler := HandleFunc(faq)
+	req, err := http.NewRequest(http.MethodGet, "", nil)
+	assert.Nil(t, err)
+	recorder := httptest.NewRecorder()
+	handler.ServeHTTP(recorder, req)
+	assert.Equal(t, http.StatusOK, recorder.Code)
+
+	var received Response
+	err = json.Unmarshal(recorder.Body.Bytes(), &received)
+	assert.Nil(t, err)
+	assert.Equal(t, Response{
+		NumRecordsExported: 20,
+		NumRecordsReceived: 15,
+		NumFlows:           30,
+		NumConnToCollector: 1,
+	}, received)
+
+	assert.Equal(t, received.GetTableRow(0), []string{"20", "15", "30", "1"})
+
+}