Skip to content

Commit

Permalink
address comment
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 29, 2022
1 parent bed014a commit 0d26e88
Show file tree
Hide file tree
Showing 9 changed files with 91 additions and 89 deletions.
28 changes: 14 additions & 14 deletions pkg/theia/commands/policy_recommendation_delete.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,26 +27,26 @@ import (
// policyRecommendationDeleteCmd represents the policy-recommendation delete command
var policyRecommendationDeleteCmd = &cobra.Command{
Use: "delete",
Short: "Delete a policy recommendation Spark job",
Long: `Delete a policy recommendation Spark job by ID.`,
Short: "Delete a policy recommendation job",
Long: `Delete a policy recommendation job by Name.`,
Aliases: []string{"del"},
Args: cobra.RangeArgs(0, 1),
Example: `
Delete the policy recommendation job with ID e998433e-accb-4888-9fc8-06563f073e86
$ theia policy-recommendation delete e998433e-accb-4888-9fc8-06563f073e86
Delete the network policy recommendation job with Name pr-e998433e-accb-4888-9fc8-06563f073e86
$ theia policy-recommendation delete pr-e998433e-accb-4888-9fc8-06563f073e86
`,
RunE: policyRecommendationDelete,
}

func policyRecommendationDelete(cmd *cobra.Command, args []string) error {
recoID, err := cmd.Flags().GetString("id")
prName, err := cmd.Flags().GetString("name")
if err != nil {
return err
}
if recoID == "" && len(args) == 1 {
recoID = args[0]
if prName == "" && len(args) == 1 {
prName = args[0]
}
err = ParseRecommendationID(recoID)
err = ParseRecommendationName(prName)
if err != nil {
return err
}
Expand All @@ -64,14 +64,14 @@ func policyRecommendationDelete(cmd *cobra.Command, args []string) error {
err = theiaClient.Delete().
AbsPath("/apis/intelligence.theia.antrea.io/v1alpha1/").
Resource("networkpolicyrecommendations").
Name("pr-" + recoID).
Name(prName).
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)
_, err = getPolicyRecommendationByName(theiaClient, prName)
if err == nil {
return false, nil
} else {
Expand All @@ -81,16 +81,16 @@ func policyRecommendationDelete(cmd *cobra.Command, args []string) error {
if err != nil {
return err
}
fmt.Printf("Successfully deleted policy recommendation job with ID %s\n", recoID)
fmt.Printf("Successfully deleted policy recommendation job with name: %s\n", prName)
return nil
}

func init() {
policyRecommendationCmd.AddCommand(policyRecommendationDeleteCmd)
policyRecommendationDeleteCmd.Flags().StringP(
"id",
"i",
"name",
"",
"ID of the policy recommendation Spark job.",
"",
"Name of the policy recommendation job.",
)
}
8 changes: 4 additions & 4 deletions pkg/theia/commands/policy_recommendation_delete_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ import (
)

func TestPolicyRecommendationDelete(t *testing.T) {
sparkAppID := "e292395c-3de1-11ed-b878-0242ac120002"
nprName := "pr-e292395c-3de1-11ed-b878-0242ac120002"
testCases := []struct {
name string
testServer *httptest.Server
Expand All @@ -40,7 +40,7 @@ func TestPolicyRecommendationDelete(t *testing.T) {
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):
case fmt.Sprintf("/apis/intelligence.theia.antrea.io/v1alpha1/networkpolicyrecommendations/%s", nprName):
if r.Method == "DELETE" {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
Expand All @@ -55,7 +55,7 @@ func TestPolicyRecommendationDelete(t *testing.T) {
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):
case fmt.Sprintf("/apis/intelligence.theia.antrea.io/v1alpha1/networkpolicyrecommendations/%s", nprName):
http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
}
})),
Expand All @@ -75,7 +75,7 @@ func TestPolicyRecommendationDelete(t *testing.T) {
SetupTheiaClientAndConnection = oldFunc
}()
cmd := new(cobra.Command)
cmd.Flags().String("id", sparkAppID, "")
cmd.Flags().String("name", nprName, "")
cmd.Flags().Bool("use-cluster-ip", true, "")
err := policyRecommendationDelete(cmd, []string{})
if tt.expectedErrorMsg == "" {
Expand Down
10 changes: 5 additions & 5 deletions pkg/theia/commands/policy_recommendation_list.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,11 @@ import (
// policyRecommendationListCmd represents the policy-recommendation list command
var policyRecommendationListCmd = &cobra.Command{
Use: "list",
Short: "List all policy recommendation Spark jobs",
Long: `List all policy recommendation Spark jobs with name, creation time and status.`,
Short: "List all policy recommendation jobs",
Long: `List all policy recommendation jobs with name, creation time and status.`,
Aliases: []string{"ls"},
Example: `
List all policy recommendation Spark jobs
List all policy recommendation jobs
$ theia policy-recommendation list
`,
RunE: policyRecommendationList,
Expand Down Expand Up @@ -62,7 +62,7 @@ func policyRecommendationList(cmd *cobra.Command, args []string) error {
}

sparkApplicationTable := [][]string{
{"CreationTime", "CompletionTime", "ID", "Status"},
{"CreationTime", "CompletionTime", "Name", "Status"},
}
for _, npr := range nprList.Items {
if npr.Status.SparkApplication == "" {
Expand All @@ -72,7 +72,7 @@ func policyRecommendationList(cmd *cobra.Command, args []string) error {
[]string{
FormatTimestamp(npr.Status.CreationTimestamp.Time),
FormatTimestamp(npr.Status.CompletionTimestamp.Time),
npr.Status.SparkApplication,
npr.Name,
npr.Status.State,
})
}
Expand Down
36 changes: 17 additions & 19 deletions pkg/theia/commands/policy_recommendation_retrieve.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,32 +24,30 @@ import (
// policyRecommendationRetrieveCmd represents the policy-recommendation retrieve command
var policyRecommendationRetrieveCmd = &cobra.Command{
Use: "retrieve",
Short: "Get the recommendation result of a policy recommendation Spark job",
Long: `Get the recommendation result of a policy recommendation Spark job by ID.
Short: "Get the recommendation result of a policy recommendation job",
Long: `Get the recommendation result of a policy recommendation job by name.
It will return the recommended NetworkPolicies described in yaml.`,
Args: cobra.RangeArgs(0, 1),
Example: `
Get the recommendation result with job ID e998433e-accb-4888-9fc8-06563f073e86
$ theia policy-recommendation retrieve --id e998433e-accb-4888-9fc8-06563f073e86
Get the recommendation result with job name pr-e998433e-accb-4888-9fc8-06563f073e86
$ theia policy-recommendation retrieve --name pr-e998433e-accb-4888-9fc8-06563f073e86
Or
$ theia policy-recommendation retrieve e998433e-accb-4888-9fc8-06563f073e86
Use a customized ClickHouse endpoint when connecting to ClickHouse to getting the result
$ theia policy-recommendation retrieve e998433e-accb-4888-9fc8-06563f073e86 --clickhouse-endpoint 10.10.1.1
Use Service ClusterIP when connecting to ClickHouse to getting the result
$ theia policy-recommendation retrieve e998433e-accb-4888-9fc8-06563f073e86 --use-cluster-ip
$ theia policy-recommendation retrieve pr-e998433e-accb-4888-9fc8-06563f073e86
Use Service ClusterIP when getting the result
$ theia policy-recommendation retrieve pr-e998433e-accb-4888-9fc8-06563f073e86 --use-cluster-ip
Save the recommendation result to file
$ theia policy-recommendation retrieve e998433e-accb-4888-9fc8-06563f073e86 --use-cluster-ip --file output.yaml
$ theia policy-recommendation retrieve pr-e998433e-accb-4888-9fc8-06563f073e86 --use-cluster-ip --file output.yaml
`,
RunE: policyRecommendationRetrieve,
}

func init() {
policyRecommendationCmd.AddCommand(policyRecommendationRetrieveCmd)
policyRecommendationRetrieveCmd.Flags().StringP(
"id",
"i",
"name",
"",
"ID of the policy recommendation Spark job.",
"",
"Name of the policy recommendation job.",
)
policyRecommendationRetrieveCmd.Flags().StringP(
"file",
Expand All @@ -60,14 +58,14 @@ func init() {
}

func policyRecommendationRetrieve(cmd *cobra.Command, args []string) error {
recoID, err := cmd.Flags().GetString("id")
prName, err := cmd.Flags().GetString("name")
if err != nil {
return err
}
if recoID == "" && len(args) == 1 {
recoID = args[0]
if prName == "" && len(args) == 1 {
prName = args[0]
}
err = ParseRecommendationID(recoID)
err = ParseRecommendationName(prName)
if err != nil {
return err
}
Expand All @@ -86,9 +84,9 @@ func policyRecommendationRetrieve(cmd *cobra.Command, args []string) error {
if pf != nil {
defer pf.Stop()
}
npr, err := getPolicyRecommendationByRecommendationID(theiaClient, recoID)
npr, err := getPolicyRecommendationByName(theiaClient, prName)
if err != nil {
return fmt.Errorf("error when getting policy recommendation job by using Recommendation ID: %v", err)
return fmt.Errorf("error when getting policy recommendation job by using job name: %v", err)
}
if filePath != "" {
if err := os.WriteFile(filePath, []byte(npr.Status.RecommendationOutcome), 0600); err != nil {
Expand Down
18 changes: 9 additions & 9 deletions pkg/theia/commands/policy_recommendation_retrieve_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,20 +35,20 @@ import (
)

func TestPolicyRecommendationRetrieve(t *testing.T) {
sparkAppID := "e292395c-3de1-11ed-b878-0242ac120002"
nprName := "pr-e292395c-3de1-11ed-b878-0242ac120002"
testCases := []struct {
name string
testServer *httptest.Server
expectedMsg []string
expectedErrorMsg string
sparkID string
nprName string
filePath 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):
case fmt.Sprintf("/apis/intelligence.theia.antrea.io/v1alpha1/networkpolicyrecommendations/%s", nprName):
npr := &intelligence.NetworkPolicyRecommendation{
Status: intelligence.NetworkPolicyRecommendationStatus{
RecommendationOutcome: "testOutcome",
Expand All @@ -59,15 +59,15 @@ func TestPolicyRecommendationRetrieve(t *testing.T) {
json.NewEncoder(w).Encode(npr)
}
})),
sparkID: "e292395c-3de1-11ed-b878-0242ac120002",
nprName: "pr-e292395c-3de1-11ed-b878-0242ac120002",
expectedMsg: []string{"testOutcome"},
expectedErrorMsg: "",
},
{
name: "Valid case with filePath",
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):
case fmt.Sprintf("/apis/intelligence.theia.antrea.io/v1alpha1/networkpolicyrecommendations/%s", nprName):
npr := &intelligence.NetworkPolicyRecommendation{
Status: intelligence.NetworkPolicyRecommendationStatus{
RecommendationOutcome: "testOutcome",
Expand All @@ -78,7 +78,7 @@ func TestPolicyRecommendationRetrieve(t *testing.T) {
json.NewEncoder(w).Encode(npr)
}
})),
sparkID: "e292395c-3de1-11ed-b878-0242ac120002",
nprName: "pr-e292395c-3de1-11ed-b878-0242ac120002",
expectedMsg: []string{"testOutcome"},
expectedErrorMsg: "",
filePath: "/tmp/testResult",
Expand All @@ -87,11 +87,11 @@ func TestPolicyRecommendationRetrieve(t *testing.T) {
name: "NetworkPolicyRecommendation 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):
case fmt.Sprintf("/apis/intelligence.theia.antrea.io/v1alpha1/networkpolicyrecommendations/%s", nprName):
http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
}
})),
sparkID: "e292395c-3de1-11ed-b878-0242ac120001",
nprName: "pr-e292395c-3de1-11ed-b878-0242ac120001",
expectedMsg: []string{},
expectedErrorMsg: "error when getting policy recommendation job",
},
Expand All @@ -111,7 +111,7 @@ func TestPolicyRecommendationRetrieve(t *testing.T) {
cmd := new(cobra.Command)
cmd.Flags().Bool("use-cluster-ip", true, "")
cmd.Flags().String("file", tt.filePath, "")
cmd.Flags().String("id", tt.sparkID, "")
cmd.Flags().String("name", tt.nprName, "")

orig := os.Stdout
r, w, _ := os.Pipe()
Expand Down
16 changes: 8 additions & 8 deletions pkg/theia/commands/policy_recommendation_run.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,16 +35,16 @@ import (
// policyRecommendationRunCmd represents the policy recommendation run command
var policyRecommendationRunCmd = &cobra.Command{
Use: "run",
Short: "Run a new policy recommendation Spark job",
Long: `Run a new policy recommendation Spark job.
Short: "Run a new policy recommendation job",
Long: `Run a new policy recommendation job.
Must finish the deployment of Theia first`,
Example: `Run a policy recommendation Spark job with default configuration
Example: `Run a policy recommendation job with default configuration
$ theia policy-recommendation run
Run an initial policy recommendation Spark job with policy type anp-deny-applied and limit on last 10k flow records
Run an initial policy recommendation job with policy type anp-deny-applied and limit on last 10k flow records
$ theia policy-recommendation run --type initial --policy-type anp-deny-applied --limit 10000
Run an initial policy recommendation Spark job with policy type anp-deny-applied and limit on flow records from 2022-01-01 00:00:00 to 2022-01-31 23:59:59.
Run an initial policy recommendation job with policy type anp-deny-applied and limit on flow records from 2022-01-01 00:00:00 to 2022-01-31 23:59:59.
$ theia policy-recommendation run --type initial --policy-type anp-deny-applied --start-time '2022-01-01 00:00:00' --end-time '2022-01-31 23:59:59'
Run a policy recommendation Spark job with default configuration but doesn't recommend toServices ANPs
Run a policy recommendation job with default configuration but doesn't recommend toServices ANPs
$ theia policy-recommendation run --to-services=false
`,
RunE: policyRecommendationRun,
Expand Down Expand Up @@ -222,9 +222,9 @@ be a list of namespace string, for example: '["kube-system","flow-aggregator","f
if waitFlag {
var npr intelligence.NetworkPolicyRecommendation
err = wait.Poll(config.StatusCheckPollInterval, config.StatusCheckPollTimeout, func() (bool, error) {
npr, err = getPolicyRecommendationByRecommendationID(theiaClient, recoID)
npr, err = getPolicyRecommendationByName(theiaClient, recoID)
if err != nil {
return false, fmt.Errorf("error when getting policy recommendation job by using Recommendation ID: %v", err)
return false, fmt.Errorf("error when getting policy recommendation job by using job name: %v", err)
}
state := npr.Status.State
if state == "COMPLETED" {
Expand Down
32 changes: 16 additions & 16 deletions pkg/theia/commands/policy_recommendation_status.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,40 +23,40 @@ import (
// policyRecommendationStatusCmd represents the policy-recommendation status command
var policyRecommendationStatusCmd = &cobra.Command{
Use: "status",
Short: "Check the status of a policy recommendation Spark job",
Long: `Check the current status of a policy recommendation Spark job by ID.
Short: "Check the status of a policy recommendation job",
Long: `Check the current status of a policy recommendation job by name.
It will return the status of this Spark application like SUBMITTED, RUNNING, COMPLETED, or FAILED.`,
Args: cobra.RangeArgs(0, 1),
Example: `
Check the current status of job with ID e998433e-accb-4888-9fc8-06563f073e86
$ theia policy-recommendation status --id e998433e-accb-4888-9fc8-06563f073e86
Check the current status of job with name pr-e998433e-accb-4888-9fc8-06563f073e86
$ theia policy-recommendation status --name pr-e998433e-accb-4888-9fc8-06563f073e86
Or
$ theia policy-recommendation status e998433e-accb-4888-9fc8-06563f073e86
Use Service ClusterIP when checking the current status of job with ID e998433e-accb-4888-9fc8-06563f073e86
$ theia policy-recommendation status e998433e-accb-4888-9fc8-06563f073e86 --use-cluster-ip
$ theia policy-recommendation status pr-e998433e-accb-4888-9fc8-06563f073e86
Use Service ClusterIP when checking the current status of job with name pr-e998433e-accb-4888-9fc8-06563f073e86
$ theia policy-recommendation status pr-e998433e-accb-4888-9fc8-06563f073e86 --use-cluster-ip
`,
RunE: policyRecommendationStatus,
}

func init() {
policyRecommendationCmd.AddCommand(policyRecommendationStatusCmd)
policyRecommendationStatusCmd.Flags().StringP(
"id",
"i",
"name",
"",
"ID of the policy recommendation Spark job.",
"",
"Name of the policy recommendation job.",
)
}

func policyRecommendationStatus(cmd *cobra.Command, args []string) error {
recoID, err := cmd.Flags().GetString("id")
prName, err := cmd.Flags().GetString("name")
if err != nil {
return err
}
if recoID == "" && len(args) == 1 {
recoID = args[0]
if prName == "" && len(args) == 1 {
prName = args[0]
}
err = ParseRecommendationID(recoID)
err = ParseRecommendationName(prName)
if err != nil {
return err
}
Expand All @@ -71,9 +71,9 @@ func policyRecommendationStatus(cmd *cobra.Command, args []string) error {
if pf != nil {
defer pf.Stop()
}
npr, err := getPolicyRecommendationByRecommendationID(theiaClient, recoID)
npr, err := getPolicyRecommendationByName(theiaClient, prName)
if err != nil {
return fmt.Errorf("error when getting policy recommendation job by using Recommendation ID: %v", err)
return fmt.Errorf("error when getting policy recommendation job by using job name: %v", err)
}
state := npr.Status.State
if state == "RUNNING" {
Expand Down
Loading

0 comments on commit 0d26e88

Please sign in to comment.