add template fields of stats in flow aggregator (#1676)

antrea-io · Dec 19, 2020 · d2d50a1 · d2d50a1
1 parent 50710ac
commit d2d50a1
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/pkg/flowaggregator/flowaggregator.go b/pkg/flowaggregator/flowaggregator.go
@@ -340,6 +340,22 @@ func (fa *flowAggregator) sendTemplateSet(templateSet ipfix.IPFIXSet) (int, erro
 		ie := ipfixentities.NewInfoElementWithValue(element, nil)
 		elements = append(elements, ie)
 	}
+	for _, ie := range antreaSourceStatsElementList {
+		element, err := fa.registry.GetInfoElement(ie, ipfixregistry.AntreaEnterpriseID)
+		if err != nil {
+			return 0, fmt.Errorf("%s not present. returned error: %v", ie, err)
+		}
+		ie := ipfixentities.NewInfoElementWithValue(element, nil)
+		elements = append(elements, ie)
+	}
+	for _, ie := range antreaDestinationStatsElementList {
+		element, err := fa.registry.GetInfoElement(ie, ipfixregistry.AntreaEnterpriseID)
+		if err != nil {
+			return 0, fmt.Errorf("%s not present. returned error: %v", ie, err)
+		}
+		ie := ipfixentities.NewInfoElementWithValue(element, nil)
+		elements = append(elements, ie)
+	}
 	err := templateSet.AddRecord(elements, fa.templateID)
 	if err != nil {
 		return 0, fmt.Errorf("error when adding record to set, error: %v", err)

diff --git a/pkg/flowaggregator/flowaggregator_test.go b/pkg/flowaggregator/flowaggregator_test.go
@@ -73,7 +73,15 @@ func TestFlowAggregator_sendTemplateSet(t *testing.T) {
 		elemList = append(elemList, ipfixentities.NewInfoElementWithValue(ipfixentities.NewInfoElement(ie, 0, 0, ipfixregistry.IANAEnterpriseID, 0), nil))
 		mockIPFIXRegistry.EXPECT().GetInfoElement(ie, ipfixregistry.IANAEnterpriseID).Return(elemList[i+len(ianaInfoElements)+len(ianaReverseInfoElements)+len(antreaInfoElements)].Element, nil)
 	}
+	for i, ie := range antreaSourceStatsElementList {
+		elemList = append(elemList, ipfixentities.NewInfoElementWithValue(ipfixentities.NewInfoElement(ie, 0, 0, ipfixregistry.AntreaEnterpriseID, 0), nil))
+		mockIPFIXRegistry.EXPECT().GetInfoElement(ie, ipfixregistry.AntreaEnterpriseID).Return(elemList[i+len(ianaInfoElements)+len(ianaReverseInfoElements)+len(antreaInfoElements)+len(aggregatorElements)].Element, nil)
+	}
 
+	for i, ie := range antreaDestinationStatsElementList {
+		elemList = append(elemList, ipfixentities.NewInfoElementWithValue(ipfixentities.NewInfoElement(ie, 0, 0, ipfixregistry.AntreaEnterpriseID, 0), nil))
+		mockIPFIXRegistry.EXPECT().GetInfoElement(ie, ipfixregistry.AntreaEnterpriseID).Return(elemList[i+len(ianaInfoElements)+len(ianaReverseInfoElements)+len(antreaInfoElements)+len(aggregatorElements)+len(antreaSourceStatsElementList)].Element, nil)
+	}
 	var tempSet ipfixentities.Set
 	mockTempSet.EXPECT().AddRecord(elemList, testTemplateID).Return(nil)
 	mockTempSet.EXPECT().GetSet().Return(tempSet)