wip

antrea-io · Jan 28, 2021 · da50f9a · da50f9a
1 parent 2a8ecf9
commit da50f9a
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 15 deletions.
diff --git a/pkg/agent/stats/collector_test.go b/pkg/agent/stats/collector_test.go
@@ -62,7 +62,7 @@ func TestCollect(t *testing.T) {
 	tests := []struct {
 		name                    string
 		ruleStats               map[uint32]*agenttypes.RuleMetric
-		ofIDToPolicyMap         map[uint32]*cpv1beta.NetworkPolicyReference
+		ofIDToPolicyMap         map[uint32]*agenttypes.PolicyRule
 		expectedStatsCollection *statsCollection
 	}{
 		{
@@ -84,10 +84,10 @@ func TestCollect(t *testing.T) {
 					Sessions: 3,
 				},
 			},
-			ofIDToPolicyMap: map[uint32]*cpv1beta.NetworkPolicyReference{
-				1: &np1,
-				2: &np1,
-				3: &np2,
+			ofIDToPolicyMap: map[uint32]*agenttypes.PolicyRule{
+				1: {PolicyRef: &np1},
+				2: {PolicyRef: &np1},
+				3: {PolicyRef: &np2},
 			},
 			expectedStatsCollection: &statsCollection{
 				networkPolicyStats: map[types.UID]*statsv1alpha1.TrafficStats{
@@ -125,10 +125,10 @@ func TestCollect(t *testing.T) {
 					Sessions: 3,
 				},
 			},
-			ofIDToPolicyMap: map[uint32]*cpv1beta.NetworkPolicyReference{
-				1: &np1,
-				2: &acnp1,
-				3: &anp1,
+			ofIDToPolicyMap: map[uint32]*agenttypes.PolicyRule{
+				1: {PolicyRef: &np1},
+				2: {Name: "rule1", PolicyRef: &acnp1},
+				3: {Name: "rule2", PolicyRef: &anp1},
 			},
 			expectedStatsCollection: &statsCollection{
 				networkPolicyStats: map[types.UID]*statsv1alpha1.TrafficStats{
@@ -172,8 +172,8 @@ func TestCollect(t *testing.T) {
 					Sessions: 1,
 				},
 			},
-			ofIDToPolicyMap: map[uint32]*cpv1beta.NetworkPolicyReference{
-				1: &np1,
+			ofIDToPolicyMap: map[uint32]*agenttypes.PolicyRule{
+				1: {PolicyRef: &np1},
 				2: nil,
 			},
 			expectedStatsCollection: &statsCollection{
@@ -195,7 +195,7 @@ func TestCollect(t *testing.T) {
 			npQuerier := queriertest.NewMockAgentNetworkPolicyInfoQuerier(ctrl)
 			ofClient.EXPECT().NetworkPolicyMetrics().Return(tt.ruleStats).Times(1)
 			for ofID, policy := range tt.ofIDToPolicyMap {
-				npQuerier.EXPECT().GetNetworkPolicyByRuleFlowID(ofID).Return(policy)
+				npQuerier.EXPECT().GetRuleByFlowID(ofID).Return(policy)
 			}
 
 			m := &Collector{ofClient: ofClient, networkPolicyQuerier: npQuerier}

diff --git a/pkg/controller/stats/aggregator_test.go b/pkg/controller/stats/aggregator_test.go
@@ -95,6 +95,11 @@ func TestAggregatorCollectListGet(t *testing.T) {
 									Packets:  5,
 									Sessions: 2,
 								},
+								"rule3": {
+									Bytes:    22,
+									Packets:  52,
+									Sessions: 22,
+								},
 							},
 						},
 					},
@@ -184,9 +189,27 @@ func TestAggregatorCollectListGet(t *testing.T) {
 						Name: cnp1.Name,
 					},
 					TrafficStats: statsv1alpha1.TrafficStats{
-						Bytes:    40,
-						Packets:  13,
-						Sessions: 7,
+						Bytes:    62,
+						Packets:  65,
+						Sessions: 29,
+					},
+					RuleTrafficStats: []statsv1alpha1.RuleTrafficStats{
+						{
+							Name: "rule1",
+							TrafficStats: statsv1alpha1.TrafficStats{
+								Bytes:    20,
+								Packets:  5,
+								Sessions: 2,
+							},
+						},
+						{
+							Name: "rule3",
+							TrafficStats: statsv1alpha1.TrafficStats{
+								Bytes:    42,
+								Packets:  60,
+								Sessions: 27,
+							},
+						},
 					},
 				},
 				{
@@ -198,6 +221,7 @@ func TestAggregatorCollectListGet(t *testing.T) {
 						Packets:  0,
 						Sessions: 0,
 					},
+					RuleTrafficStats: nil,
 				},
 			},
 			expectedAntreaNetworkPolicyStats: []statsv1alpha1.AntreaNetworkPolicyStats{
@@ -211,6 +235,24 @@ func TestAggregatorCollectListGet(t *testing.T) {
 						Packets:  15,
 						Sessions: 7,
 					},
+					RuleTrafficStats: []statsv1alpha1.RuleTrafficStats{
+						{
+							Name: "rule2",
+							TrafficStats: statsv1alpha1.TrafficStats{
+								Bytes:    20,
+								Packets:  5,
+								Sessions: 2,
+							},
+						},
+						{
+							Name: "rule4",
+							TrafficStats: statsv1alpha1.TrafficStats{
+								Bytes:    100,
+								Packets:  10,
+								Sessions: 5,
+							},
+						},
+					},
 				},
 				{
 					ObjectMeta: metav1.ObjectMeta{
@@ -345,12 +387,14 @@ func TestAggregatorCollectListGet(t *testing.T) {
 				actualStats, exists := a.GetAntreaClusterNetworkPolicyStats(Stats.Name)
 				require.True(t, exists)
 				require.Equal(t, Stats.TrafficStats, actualStats.TrafficStats)
+				require.ElementsMatch(t, Stats.RuleTrafficStats, actualStats.RuleTrafficStats)
 			}
 			assert.Equal(t, len(tt.expectedAntreaNetworkPolicyStats), len(a.ListAntreaNetworkPolicyStats("")))
 			for _, Stats := range tt.expectedAntreaNetworkPolicyStats {
 				actualStats, exists := a.GetAntreaNetworkPolicyStats(Stats.Namespace, Stats.Name)
 				require.True(t, exists)
 				require.Equal(t, Stats.TrafficStats, actualStats.TrafficStats)
+				require.ElementsMatch(t, Stats.RuleTrafficStats, actualStats.RuleTrafficStats)
 			}
 		})
 	}