Refactor UT and add more testcases

Signed-off-by: Yongming Ding <dyongming@vmware.com>

Makefile

@@ -74,7 +74,7 @@ add-copyright:
 .linux-test-unit:
 	@echo
 	@echo "==> Running unit tests <=="
-	$(GO) test -race -covermode=atomic -cover antrea.io/theia/plugins/... 
+	$(GO) test -race -covermode=atomic -cover antrea.io/theia/plugins/... antrea.io/theia/pkg/... 
 
 .PHONY: tidy
 tidy:

pkg/theia/commands/policy_recommendation_retrieve.go

@@ -144,7 +144,7 @@ func getPolicyRecommendationResult(clientset kubernetes.Interface, kubeconfig st
 	}
 	recoResult, err = getResultFromClickHouse(connect, recoID)
 	if err != nil {
-		return "", fmt.Errorf("error when connecting to ClickHouse, %v", err)
+		return "", fmt.Errorf("error when getting result from ClickHouse, %v", err)
 	}
 	if filePath != "" {
 		if err := ioutil.WriteFile(filePath, []byte(recoResult), 0600); err != nil {

pkg/theia/commands/policy_recommendation_retrieve_test.go

@@ -25,33 +25,122 @@ import (
 )
 
 func TestGetClickHouseSecret(t *testing.T) {
-	var fakeClientset = fake.NewSimpleClientset(
-		&v1.Secret{
-			ObjectMeta: metav1.ObjectMeta{
-				Name:      "clickhouse-secret",
-				Namespace: flowVisibilityNS,
-			},
-			Data: map[string][]byte{
-				"password": []byte("clickhouse_operator_password"),
-				"username": []byte("clickhouse_operator"),
-			},
+	testCases := []struct {
+		name             string
+		fakeClientset    *fake.Clientset
+		expectedUsername string
+		expectedPassword string
+		expectedErrorMsg string
+	}{
+		{
+			name: "valid case",
+			fakeClientset: fake.NewSimpleClientset(
+				&v1.Secret{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "clickhouse-secret",
+						Namespace: flowVisibilityNS,
+					},
+					Data: map[string][]byte{
+						"username": []byte("clickhouse_operator"),
+						"password": []byte("clickhouse_operator_password"),
+					},
+				},
+			),
+			expectedUsername: "clickhouse_operator",
+			expectedPassword: "clickhouse_operator_password",
+			expectedErrorMsg: "",
 		},
-	)
-	username, password, err := getClickHouseSecret(fakeClientset)
-	assert.NoError(t, err)
-	assert.Equal(t, "clickhouse_operator", string(username))
-	assert.Equal(t, "clickhouse_operator_password", string(password))
+		{
+			name:             "clickhouse secret not found",
+			fakeClientset:    fake.NewSimpleClientset(),
+			expectedUsername: "",
+			expectedPassword: "",
+			expectedErrorMsg: `error secrets "clickhouse-secret" not found when finding the ClickHouse secret, please check the deployment of ClickHouse`,
+		},
+		{
+			name: "username not found",
+			fakeClientset: fake.NewSimpleClientset(
+				&v1.Secret{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "clickhouse-secret",
+						Namespace: flowVisibilityNS,
+					},
+					Data: map[string][]byte{
+						"password": []byte("clickhouse_operator_password"),
+					},
+				},
+			),
+			expectedUsername: "",
+			expectedPassword: "",
+			expectedErrorMsg: "error when getting the ClickHouse username",
+		},
+		{
+			name: "password not found",
+			fakeClientset: fake.NewSimpleClientset(
+				&v1.Secret{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "clickhouse-secret",
+						Namespace: flowVisibilityNS,
+					},
+					Data: map[string][]byte{
+						"username": []byte("clickhouse_operator"),
+					},
+				},
+			),
+			expectedUsername: "clickhouse_operator",
+			expectedPassword: "",
+			expectedErrorMsg: "error when getting the ClickHouse password",
+		},
+	}
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			username, password, err := getClickHouseSecret(tt.fakeClientset)
+			if tt.expectedErrorMsg != "" {
+				assert.EqualErrorf(t, err, tt.expectedErrorMsg, "Error should be: %v, got: %v", tt.expectedErrorMsg, err)
+			}
+			assert.Equal(t, tt.expectedUsername, string(username))
+			assert.Equal(t, tt.expectedPassword, string(password))
+		})
+	}
 }
 
 func TestGetResultFromClickHouse(t *testing.T) {
-	recoID := "db2134ea-7169-46f8-b56d-d643d4751d1d"
-	expectedResult := "recommend-allow-acnp-kube-system-rpeal"
-	db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
-	assert.NoError(t, err)
-	defer db.Close()
-	resultRow := sqlmock.NewRows([]string{"yamls"}).AddRow(expectedResult)
-	mock.ExpectQuery("SELECT yamls FROM recommendations WHERE id = (?);").WithArgs(recoID).WillReturnRows(resultRow)
-	result, err := getResultFromClickHouse(db, recoID)
-	assert.NoError(t, err)
-	assert.Equal(t, expectedResult, result)
+	testCases := []struct {
+		name             string
+		recommendationID string
+		expectedResult   string
+		expectedErrorMsg string
+	}{
+		{
+			name:             "valid case",
+			recommendationID: "db2134ea-7169-46f8-b56d-d643d4751d1d",
+			expectedResult:   "recommend-allow-acnp-kube-system-rpeal",
+			expectedErrorMsg: "",
+		},
+		{
+			name:             "no result given recommendation ID",
+			recommendationID: "db2134ea-7169",
+			expectedResult:   "",
+			expectedErrorMsg: "failed to get recommendation result with id db2134ea-7169: sql: no rows in result set",
+		},
+	}
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
+			assert.NoError(t, err)
+			defer db.Close()
+			resultRow := &sqlmock.Rows{}
+			if tt.expectedResult != "" {
+				resultRow = sqlmock.NewRows([]string{"yamls"}).AddRow(tt.expectedResult)
+			}
+			mock.ExpectQuery("SELECT yamls FROM recommendations WHERE id = (?);").WithArgs(tt.recommendationID).WillReturnRows(resultRow)
+			result, err := getResultFromClickHouse(db, tt.recommendationID)
+			if tt.expectedErrorMsg != "" {
+				assert.EqualErrorf(t, err, tt.expectedErrorMsg, "Error should be: %v, got: %v", tt.expectedErrorMsg, err)
+			} else {
+				assert.NoError(t, err)
+			}
+			assert.Equal(t, tt.expectedResult, result)
+		})
+	}
 }

pkg/theia/commands/policy_recommendation_run.go

@@ -35,7 +35,7 @@ import (
 const (
 	flowVisibilityNS        = "flow-visibility"
 	k8sQuantitiesReg        = "^([+-]?[0-9.]+)([eEinumkKMGTP]*[-+]?[0-9]*)$"
-	sparkImage              = "antrea/theia-policy-recommendation:latest"
+	sparkImage              = "projects.registry.vmware.com/antrea/theia-policy-recommendation:latest"
 	sparkImagePullPolicy    = "IfNotPresent"
 	sparkAppFile            = "local:///opt/spark/work-dir/policy_recommendation_job.py"
 	sparkServiceAccount     = "policy-recommendation-spark"

pkg/theia/commands/policy_recommendation_status_test.go

@@ -27,31 +27,88 @@ import (
 
 func TestGetPolicyRecommendationProgress(t *testing.T) {
 	sparkAppID := "spark-0fa6cc19ae23439794747a306d5ad705"
-	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		switch strings.TrimSpace(r.URL.Path) {
-		case "/api/v1/applications":
-			responses := []map[string]interface{}{
-				{"id": sparkAppID},
+	testCases := []struct {
+		name             string
+		testServer       *httptest.Server
+		expectedProgress string
+		expectedErrorMsg string
+	}{
+		{
+			name: "valid case",
+			testServer: httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				switch strings.TrimSpace(r.URL.Path) {
+				case "/api/v1/applications":
+					responses := []map[string]interface{}{
+						{"id": sparkAppID},
+					}
+					w.Header().Set("Content-Type", "application/json")
+					w.WriteHeader(http.StatusOK)
+					json.NewEncoder(w).Encode(responses)
+				case fmt.Sprintf("/api/v1/applications/%s/stages", sparkAppID):
+					responses := []map[string]interface{}{
+						{"status": "COMPLETE"},
+						{"status": "COMPLETE"},
+						{"status": "SKIPPED"},
+						{"status": "PENDING"},
+						{"status": "ACTIVE"},
+					}
+					w.Header().Set("Content-Type", "application/json")
+					w.WriteHeader(http.StatusOK)
+					json.NewEncoder(w).Encode(responses)
+				}
+			})),
+			expectedProgress: ": 3/5 (60%) stages completed",
+			expectedErrorMsg: "",
+		},
+		{
+			name: "found more than one spark application",
+			testServer: httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				switch strings.TrimSpace(r.URL.Path) {
+				case "/api/v1/applications":
+					responses := []map[string]interface{}{
+						{"id": sparkAppID},
+						{"id": sparkAppID},
+					}
+					w.Header().Set("Content-Type", "application/json")
+					w.WriteHeader(http.StatusOK)
+					json.NewEncoder(w).Encode(responses)
+				}
+			})),
+			expectedProgress: "",
+			expectedErrorMsg: "wrong Spark Application number, expected 1, got 2",
+		},
+		{
+			name: "no spark application stage found",
+			testServer: httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				switch strings.TrimSpace(r.URL.Path) {
+				case "/api/v1/applications":
+					responses := []map[string]interface{}{
+						{"id": sparkAppID},
+					}
+					w.Header().Set("Content-Type", "application/json")
+					w.WriteHeader(http.StatusOK)
+					json.NewEncoder(w).Encode(responses)
+				case fmt.Sprintf("/api/v1/applications/%s/stages", sparkAppID):
+					responses := []map[string]interface{}{}
+					w.Header().Set("Content-Type", "application/json")
+					w.WriteHeader(http.StatusOK)
+					json.NewEncoder(w).Encode(responses)
+				}
+			})),
+			expectedProgress: "",
+			expectedErrorMsg: "wrong Spark Application stages number, expected at least 1, got 0",
+		},
+	}
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			defer tt.testServer.Close()
+			progress, err := getPolicyRecommendationProgress(tt.testServer.URL)
+			if tt.expectedErrorMsg != "" {
+				assert.EqualErrorf(t, err, tt.expectedErrorMsg, "Error should be: %v, got: %v", tt.expectedErrorMsg, err)
+			} else {
+				assert.NoError(t, err)
 			}
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusOK)
-			json.NewEncoder(w).Encode(responses)
-		case fmt.Sprintf("/api/v1/applications/%s/stages", sparkAppID):
-			responses := []map[string]interface{}{
-				{"status": "COMPLETE"},
-				{"status": "COMPLETE"},
-				{"status": "SKIPPED"},
-				{"status": "PENDING"},
-				{"status": "ACTIVE"},
-			}
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusOK)
-			json.NewEncoder(w).Encode(responses)
-		}
-	}))
-	defer server.Close()
-	expectedProgress := ": 3/5 (60%) stages completed"
-	progress, err := getPolicyRecommendationProgress(server.URL)
-	assert.NoError(t, err)
-	assert.Equal(t, expectedProgress, progress)
+			assert.Equal(t, tt.expectedProgress, progress)
+		})
+	}
 }