Skip to content

Commit

Permalink
Add unit test
Browse files Browse the repository at this point in the history
Signed-off-by: Yun-Tang Hsu <hsuy@vmware.com>
  • Loading branch information
Yun-Tang Hsu committed Sep 27, 2022
1 parent c9cab76 commit 153764b
Show file tree
Hide file tree
Showing 13 changed files with 1,047 additions and 545 deletions.
2 changes: 1 addition & 1 deletion pkg/theia/commands/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ const (
StatusCheckPollTimeout = 60 * time.Minute
CAConfigMapName = "theia-ca"
CAConfigMapKey = "ca.crt"
TokenSecretName = "theia-cli-account-token"
TokenName = "theia-cli-account-token"
ServiceAccountTokenKey = "token"
TheiaManagerServiceName = "theia-manager"
)
96 changes: 47 additions & 49 deletions pkg/theia/commands/policy_recommendation_delete.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,56 +35,54 @@ var policyRecommendationDeleteCmd = &cobra.Command{
Delete the policy recommendation job with ID e998433e-accb-4888-9fc8-06563f073e86
$ theia policy-recommendation delete e998433e-accb-4888-9fc8-06563f073e86
`,
RunE: func(cmd *cobra.Command, args []string) error {
recoID, err := cmd.Flags().GetString("id")
if err != nil {
return err
}
if recoID == "" && len(args) == 1 {
recoID = args[0]
}
err = ParseRecommendationID(recoID)
if err != nil {
return err
}
useClusterIP, err := cmd.Flags().GetBool("use-cluster-ip")
if err != nil {
return err
}
theiaClient, pf, err := SetupTheiaClientAndConnection(cmd, useClusterIP)
if err != nil {
return fmt.Errorf("couldn't setup theia client, %v", err)
}
if pf != nil {
defer pf.Stop()
}
npr, err := getPolicyRecommendationByRecommendationID(theiaClient, recoID)
if err != nil {
return fmt.Errorf("error when getting policy recommendation job by using Recommendation ID: %v", err)
}
err = theiaClient.CoreV1().RESTClient().Delete().
AbsPath("/apis/intelligence.theia.antrea.io/v1alpha1/").
Resource("networkpolicyrecommendations").
Name(npr.Name).
Do(context.TODO()).
Error()
if err != nil {
return fmt.Errorf("error when deleting policy recommendation job in Theia Manager: %v", err)
}
err = wait.Poll(config.StatusCheckPollInterval, config.StatusCheckPollTimeout, func() (bool, error) {
npr, err = getPolicyRecommendationByRecommendationID(theiaClient, recoID)
if err == nil {
return false, nil
} else {
return true, nil
}
})
if err != nil {
return err
RunE: policyRecommendationDelete,
}

func policyRecommendationDelete(cmd *cobra.Command, args []string) error {
recoID, err := cmd.Flags().GetString("id")
if err != nil {
return err
}
if recoID == "" && len(args) == 1 {
recoID = args[0]
}
err = ParseRecommendationID(recoID)
if err != nil {
return err
}
useClusterIP, err := cmd.Flags().GetBool("use-cluster-ip")
if err != nil {
return err
}
theiaClient, pf, err := SetupTheiaClientAndConnection(cmd, useClusterIP)
if err != nil {
return fmt.Errorf("couldn't setup theia client, %v", err)
}
if pf != nil {
defer pf.Stop()
}
err = theiaClient.Delete().
AbsPath("/apis/intelligence.theia.antrea.io/v1alpha1/").
Resource("networkpolicyrecommendations").
Name("pr-" + recoID).
Do(context.TODO()).
Error()
if err != nil {
return fmt.Errorf("error when deleting policy recommendation job in Theia Manager: %v", err)
}
err = wait.Poll(config.StatusCheckPollInterval, config.StatusCheckPollTimeout, func() (bool, error) {
_, err = getPolicyRecommendationByRecommendationID(theiaClient, recoID)
if err == nil {
return false, nil
} else {
return true, nil
}
fmt.Printf("Successfully deleted policy recommendation job with ID %s\n", recoID)
return nil
},
})
if err != nil {
return err
}
fmt.Printf("Successfully deleted policy recommendation job with ID %s\n", recoID)
return nil
}

func init() {
Expand Down
89 changes: 89 additions & 0 deletions pkg/theia/commands/policy_recommendation_delete_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
// 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 commands

import (
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"

"github.com/spf13/cobra"
"github.com/stretchr/testify/assert"
"k8s.io/client-go/kubernetes"
restclient "k8s.io/client-go/rest"

"antrea.io/theia/pkg/theia/portforwarder"
)

func TestPolicyRecommendationDelete(t *testing.T) {
sparkAppID := "e292395c-3de1-11ed-b878-0242ac120002"
testCases := []struct {
name string
testServer *httptest.Server
expectedErrorMsg string
}{
{
name: "valid case",
testServer: httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch strings.TrimSpace(r.URL.Path) {
case fmt.Sprintf("/apis/intelligence.theia.antrea.io/v1alpha1/networkpolicyrecommendations/pr-%s", sparkAppID):
if r.Method == "DELETE" {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
} else {
http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
}
}
})),
expectedErrorMsg: "",
},
{
name: "SparkApplication not found",
testServer: httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch strings.TrimSpace(r.URL.Path) {
case fmt.Sprintf("/apis/intelligence.theia.antrea.io/v1alpha1/networkpolicyrecommendations/pr-%s", sparkAppID):
http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
}
})),
expectedErrorMsg: "error when deleting policy recommendation job",
},
}
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
defer tt.testServer.Close()
oldFunc := SetupTheiaClientAndConnection
SetupTheiaClientAndConnection = func(cmd *cobra.Command, useClusterIP bool) (restclient.Interface, *portforwarder.PortForwarder, error) {
clientConfig := &restclient.Config{Host: tt.testServer.URL, TLSClientConfig: restclient.TLSClientConfig{Insecure: true}}
clientset, _ := kubernetes.NewForConfig(clientConfig)
return clientset.CoreV1().RESTClient(), nil, nil
}
defer func() {
SetupTheiaClientAndConnection = oldFunc
}()
cmd := new(cobra.Command)
cmd.Flags().String("id", sparkAppID, "")
cmd.Flags().Bool("use-cluster-ip", true, "")
err := policyRecommendationDelete(cmd, []string{})
if tt.expectedErrorMsg == "" {
assert.NoError(t, err)
} else {
assert.Error(t, err)
assert.Contains(t, err.Error(), tt.expectedErrorMsg)
}
})
}
}
80 changes: 41 additions & 39 deletions pkg/theia/commands/policy_recommendation_list.go
Original file line number Diff line number Diff line change
Expand Up @@ -33,47 +33,49 @@ var policyRecommendationListCmd = &cobra.Command{
List all policy recommendation Spark jobs
$ theia policy-recommendation list
`,
RunE: func(cmd *cobra.Command, args []string) error {
useClusterIP, err := cmd.Flags().GetBool("use-cluster-ip")
if err != nil {
return err
}
theiaClient, pf, err := SetupTheiaClientAndConnection(cmd, useClusterIP)
if err != nil {
return fmt.Errorf("couldn't setup theia client, %v", err)
}
if pf != nil {
defer pf.Stop()
}
nprList := &intelligence.NetworkPolicyRecommendationList{}
err = theiaClient.CoreV1().RESTClient().Get().
AbsPath("/apis/intelligence.theia.antrea.io/v1alpha1/").
Resource("networkpolicyrecommendations").
Do(context.TODO()).Into(nprList)
if err != nil {
return fmt.Errorf("error when getting policy recommendation jobs from Theia Manager: %v", err)
}

sparkApplicationTable := [][]string{
{"CreationTime", "CompletionTime", "ID", "Status"},
}
for _, npr := range nprList.Items {
if npr.Status.SparkApplication == "" {
continue
}
sparkApplicationTable = append(sparkApplicationTable,
[]string{
FormatTimestamp(npr.Status.CreationTimestamp.Time),
FormatTimestamp(npr.Status.CompletionTimestamp.Time),
npr.Status.SparkApplication,
npr.Status.State,
})
}
TableOutput(sparkApplicationTable)
return nil
},
RunE: policyRecommendationList,
}

func init() {
policyRecommendationCmd.AddCommand(policyRecommendationListCmd)
}

func policyRecommendationList(cmd *cobra.Command, args []string) error {
useClusterIP, err := cmd.Flags().GetBool("use-cluster-ip")
if err != nil {
return err
}
theiaClient, pf, err := SetupTheiaClientAndConnection(cmd, useClusterIP)
if err != nil {
return fmt.Errorf("couldn't setup theia client, %v", err)
}
if pf != nil {
defer pf.Stop()
}
nprList := &intelligence.NetworkPolicyRecommendationList{}
err = theiaClient.Get().
AbsPath("/apis/intelligence.theia.antrea.io/v1alpha1/").
Resource("networkpolicyrecommendations").
Do(context.TODO()).Into(nprList)
if err != nil {
return fmt.Errorf("error when getting policy recommendation jobs from Theia Manager: %v", err)
}

sparkApplicationTable := [][]string{
{"CreationTime", "CompletionTime", "ID", "Status"},
}
for _, npr := range nprList.Items {
if npr.Status.SparkApplication == "" {
continue
}
sparkApplicationTable = append(sparkApplicationTable,
[]string{
FormatTimestamp(npr.Status.CreationTimestamp.Time),
FormatTimestamp(npr.Status.CompletionTimestamp.Time),
npr.Status.SparkApplication,
npr.Status.State,
})
}
TableOutput(sparkApplicationTable)
return nil
}
107 changes: 107 additions & 0 deletions pkg/theia/commands/policy_recommendation_list_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
// 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 commands

import (
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"

"github.com/spf13/cobra"
"github.com/stretchr/testify/assert"
"k8s.io/client-go/kubernetes"
restclient "k8s.io/client-go/rest"

intelligence "antrea.io/theia/pkg/apis/intelligence/v1alpha1"
"antrea.io/theia/pkg/theia/portforwarder"
)

func TestPolicyRecommendationList(t *testing.T) {
testCases := []struct {
name string
testServer *httptest.Server
expectedMsg []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 fmt.Sprintf("/apis/intelligence.theia.antrea.io/v1alpha1/networkpolicyrecommendations"):
nprList := &intelligence.NetworkPolicyRecommendationList{
Items: []intelligence.NetworkPolicyRecommendation{
{Status: intelligence.NetworkPolicyRecommendationStatus{
SparkApplication: "test1",
}},
},
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(nprList)
}
})),
expectedMsg: []string{"test1"},
expectedErrorMsg: "",
},
{
name: "NetworkPolicyRecommendationList not found",
testServer: httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch strings.TrimSpace(r.URL.Path) {
case fmt.Sprintf("/apis/intelligence.theia.antrea.io/v1alpha1/networkpolicyrecommendations"):
http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
}
})),
expectedMsg: []string{},
expectedErrorMsg: "error when getting policy recommendation jobs from Theia Manager",
},
}
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
defer tt.testServer.Close()
oldFunc := SetupTheiaClientAndConnection
SetupTheiaClientAndConnection = func(cmd *cobra.Command, useClusterIP bool) (restclient.Interface, *portforwarder.PortForwarder, error) {
clientConfig := &restclient.Config{Host: tt.testServer.URL, TLSClientConfig: restclient.TLSClientConfig{Insecure: true}}
clientset, _ := kubernetes.NewForConfig(clientConfig)
return clientset.CoreV1().RESTClient(), nil, nil
}
defer func() {
SetupTheiaClientAndConnection = oldFunc
}()
cmd := new(cobra.Command)
cmd.Flags().Bool("use-cluster-ip", true, "")

orig := os.Stdout
r, w, _ := os.Pipe()
os.Stdout = w
err := policyRecommendationList(cmd, []string{})
if tt.expectedErrorMsg == "" {
assert.NoError(t, err)
outcome := readStdout(t, r, w)
os.Stdout = orig
assert.Contains(t, outcome, "test1")
for _, msg := range tt.expectedMsg {
assert.Contains(t, outcome, msg)
}
} else {
assert.Error(t, err)
assert.Contains(t, err.Error(), tt.expectedErrorMsg)
}
})
}
}
Loading

0 comments on commit 153764b

Please sign in to comment.