Merge pull request #22 from ShohamBit/test_commands

Design changes to traceectl

cmd/event.go

@@ -3,6 +3,9 @@ package cmd
 import (
 	"context"
 
+	"github.com/ShohamBit/traceectl/pkg/client"
+	"github.com/ShohamBit/traceectl/pkg/cmd/formatter"
+	"github.com/ShohamBit/traceectl/pkg/cmd/printer"
 	pb "github.com/aquasecurity/tracee/api/v1beta1"
 	"github.com/spf13/cobra"
 )
@@ -57,7 +60,7 @@ var describeEventCmd = &cobra.Command{
 	Long:  `Retrieves the detailed definition of a specific event, including its fields, types, and other metadata.`,
 	Args:  cobra.MinimumNArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
-
+		getEventDescriptions(cmd, args)
 	},
 }
 
@@ -93,14 +96,16 @@ var runEventCmd = &cobra.Command{
 
 func enableEvents(cmd *cobra.Command, eventNames []string) {
 	// Create Tracee gRPC client
-	if err := TCS.NewServiceClient(serverInfo); err != nil {
+	var traceeClient client.ServiceClient // tracee client
+
+	if err := traceeClient.NewServiceClient(serverInfo); err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return // Exit on error
 	}
 
 	// Iterate over event names and enable each one
 	for _, eventName := range eventNames {
-		_, err := TCS.EnableEvent(context.Background(), &pb.EnableEventRequest{Name: eventName})
+		_, err := traceeClient.EnableEvent(context.Background(), &pb.EnableEventRequest{Name: eventName})
 		if err != nil {
 			cmd.PrintErrln("Error enabling event:", err)
 			continue // Continue on error with the next event
@@ -111,18 +116,45 @@ func enableEvents(cmd *cobra.Command, eventNames []string) {
 
 func disableEvents(cmd *cobra.Command, eventNames []string) {
 	// Create Tracee gRPC client
-	if err := TCS.NewServiceClient(serverInfo); err != nil {
+	var traceeClient client.ServiceClient
+	if err := traceeClient.NewServiceClient(serverInfo); err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return // Exit on error
 	}
 
 	// Iterate over event names and disable each one
 	for _, eventName := range eventNames {
-		_, err := TCS.DisableEvent(context.Background(), &pb.DisableEventRequest{Name: eventName})
+		_, err := traceeClient.DisableEvent(context.Background(), &pb.DisableEventRequest{Name: eventName})
 		if err != nil {
 			cmd.PrintErrln("Error disabling event:", err)
 			continue // Continue on error with the next event
 		}
 		cmd.Println("Disabled event:", eventName)
 	}
 }
+
+func getEventDescriptions(cmd *cobra.Command, args []string) {
+	//create service client
+	var traceeClient client.ServiceClient
+	if err := traceeClient.NewServiceClient(serverInfo); err != nil {
+		cmd.PrintErrln("Error creating client: ", err)
+	}
+	defer traceeClient.CloseConnection()
+	response, err := traceeClient.GetEventDefinitions(context.Background(), &pb.GetEventDefinitionsRequest{EventNames: args})
+
+	if err != nil {
+		cmd.PrintErrln("Error getting event definitions: ", err)
+		return
+
+	}
+	//display event definitions
+	//don't support different outputs and formats
+	format, err := formatter.New("table", "", cmd)
+	if err != nil {
+		cmd.PrintErrln("Error creating formatter: ", err)
+		return
+	}
+	//show events
+	printer.DescribeEvent(format, args, response)
+
+}

cmd/event_test.go

@@ -45,7 +45,7 @@ var eventTests = []models.TestCase{
 
 func TestEvent(t *testing.T) {
 	// Start the mock server
-	mockServer, err := mock.StartMockServiceServer()
+	mockServer, err := mock.StartMockServer()
 	if err != nil {
 		t.Fatalf("Failed to start mock server: %v", err)
 	}

cmd/plugin_test.go

@@ -19,7 +19,7 @@ var pluginTests = []models.TestCase{
 
 func TestPluginCmd(t *testing.T) {
 	// Start the mock server
-	mockServer, err := mock.StartMockServiceServer()
+	mockServer, err := mock.StartMockServer()
 	if err != nil {
 		t.Fatalf("Failed to start mock server: %v", err)
 	}

cmd/policy_test.go

@@ -19,7 +19,7 @@ var policyTests = []models.TestCase{
 
 func TestPolicyCmd(t *testing.T) {
 	// Start the mock server
-	mockServer, err := mock.StartMockServiceServer()
+	mockServer, err := mock.StartMockServer()
 	if err != nil {
 		t.Fatalf("Failed to start mock server: %v", err)
 	}

cmd/root.go

@@ -2,7 +2,6 @@ package cmd
 
 import (
 	"context"
-	"fmt"
 	"os"
 
 	"github.com/ShohamBit/traceectl/pkg/client"
@@ -12,9 +11,7 @@ import (
 )
 
 var (
-	TCS        client.ServiceClient    // tracee service client
-	TCD        client.DiagnosticClient // tracee diagnostic  client
-	serverInfo client.ServerInfo       = client.ServerInfo{
+	serverInfo client.ServerInfo = client.ServerInfo{
 		ConnectionType: client.PROTOCOL_UNIX,
 		UnixSocketPath: client.SOCKET,
 		ADDR:           client.DefaultIP + ":" + client.DefaultPort,
@@ -56,7 +53,7 @@ func init() {
 	//unix connection type flag
 	rootCmd.PersistentFlags().StringVar(&serverInfo.UnixSocketPath, "socketPath", client.SOCKET, "Path of the unix socket")
 	//tcp connection type flag
-	rootCmd.PersistentFlags().StringVarP(&serverInfo.ADDR, "server", "s", client.DefaultIP+":"+client.DefaultPort, "he address and port of the Kubernetes API server")
+	rootCmd.PersistentFlags().StringVarP(&serverInfo.ADDR, "server", "s", client.DefaultIP+":"+client.DefaultPort, "The address and port of the Kubernetes API server")
 
 }
 
@@ -65,7 +62,6 @@ var connectCmd = &cobra.Command{
 	Short: "Connect to the server",
 	Long:  "Connects to a stream and displays events in real time.",
 	Run: func(cmd *cobra.Command, args []string) {
-
 	},
 }
 var metricsCmd = &cobra.Command{
@@ -104,7 +100,6 @@ var configCmd = &cobra.Command{
 	Short: "View or modify the Tracee Daemon configuration at runtime.",
 	Long:  `View or modify the Tracee Daemon configuration at runtime.`,
 	Run: func(cmd *cobra.Command, args []string) {
-		fmt.Println("config called")
 	},
 }
 
@@ -131,16 +126,18 @@ func GetRootCmd() *cobra.Command {
 
 // displayMetrics fetches and prints Tracee metrics
 func displayMetrics(cmd *cobra.Command, _ []string) {
-
+	var traceeClient client.DiagnosticClient
 	//create service client
-	if err := TCD.NewDiagnosticClient(serverInfo); err == nil {
+	if err := traceeClient.NewDiagnosticClient(serverInfo); err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
+		return
 	}
-	defer TCD.CloseConnection()
+	defer traceeClient.CloseConnection()
 	//get metrics
-	response, err := TCD.GetMetrics(context.Background(), &pb.GetMetricsRequest{})
+	response, err := traceeClient.GetMetrics(context.Background(), &pb.GetMetricsRequest{})
 	if err != nil {
-		cmd.PrintErrln("Error getting version: ", err)
+		cmd.PrintErrln("Error getting metrics: ", err)
+		return
 	}
 
 	// Display the metrics
@@ -156,17 +153,20 @@ func displayMetrics(cmd *cobra.Command, _ []string) {
 }
 
 func displayVersion(cmd *cobra.Command, _ []string) {
-
 	//create service client
-	if err := TCS.NewServiceClient(serverInfo); err != nil {
+	var traceeClient client.ServiceClient
+	if err := traceeClient.NewServiceClient(serverInfo); err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 	}
-	defer TCS.CloseConnection()
+	defer traceeClient.CloseConnection()
 	//get version
-	response, err := TCS.GetVersion(context.Background(), &pb.GetVersionRequest{})
+	response, err := traceeClient.GetVersion(context.Background(), &pb.GetVersionRequest{})
+
 	if err != nil {
 		cmd.PrintErrln("Error getting version: ", err)
+		return
+	} else {
+		//display version
+		cmd.Println("Version: ", response.Version)
 	}
-	//display version
-	cmd.Println("Version: ", response.Version)
 }

cmd/root_test.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"bytes"
+	"strings"
 	"testing"
 	"time"
 
@@ -11,15 +12,25 @@ import (
 
 var rootTests = []models.TestCase{
 	{
-		Name:           "No root subcommand",
-		Args:           []string{"root"},
-		ExpectedOutput: rootCmd.Help(), // Update expected output
+		Name:           "version",
+		Args:           []string{"version"},
+		ExpectedOutput: mock.ExpectedVersion,
+	},
+	{
+		Name: "metrics",
+		Args: []string{"metrics"},
+		ExpectedOutput: (func() string {
+			str := mock.ExpectedMetrics.String()      // convert to string
+			str = strings.ReplaceAll(str, "  ", "\n") // add newlines
+			str = strings.ReplaceAll(str, ":", ": ")  // add spaces
+			return str
+		})(),
 	},
 }
 
 func TestRootCmd(t *testing.T) {
 	// Start the mock server
-	mockServer, err := mock.StartMockServiceServer()
+	mockServer, err := mock.StartMockServer()
 	if err != nil {
 		t.Fatalf("Failed to start mock server: %v", err)
 	}
@@ -47,8 +58,8 @@ func TestRootCmd(t *testing.T) {
 			// Validate output and error (if any)
 			output := buf.String()
 
-			if output != test.ExpectedOutput {
-				t.Errorf("Expected output: %s, got: %s", test.ExpectedOutput, output)
+			if !strings.Contains(output, test.ExpectedOutput.(string)) {
+				t.Errorf("Expected output:\n%s\ngot:\n%s", test.ExpectedOutput, output)
 			}
 		})
 	}

cmd/stream.go

@@ -1,9 +1,11 @@
 package cmd
 
 import (
+	"github.com/ShohamBit/traceectl/pkg/client"
+	"github.com/ShohamBit/traceectl/pkg/cmd/formatter"
+	"github.com/ShohamBit/traceectl/pkg/cmd/printer"
 	pb "github.com/aquasecurity/tracee/api/v1beta1"
 
-	"github.com/ShohamBit/traceectl/pkg/cmd/formatter"
 	"github.com/spf13/cobra"
 )
 
@@ -13,6 +15,7 @@ var streamCmd = &cobra.Command{
 	Use:   "stream [policies...]",
 	Short: "Stream events from tracee",
 	Long: `Stream Management:
+	- traceectl stream [POLICIES...] -  stream event directly from tracee
   	- traceectl stream create --name <stream_name> [--destination <destination>] [--format <format>] [--fields <fields>] [--parse-data] [--filter <filter>]
   	- traceectl stream describe <stream_name>
   	- traceectl stream list
@@ -117,69 +120,33 @@ var resumeStreamCmd = &cobra.Command{
 	},
 }
 
+// stream events directly from tracee
 func stream(cmd *cobra.Command, args []string) {
 	// Create service client
-	err := TCS.NewServiceClient(serverInfo)
+	var traceeClient client.ServiceClient
+	err := traceeClient.NewServiceClient(serverInfo)
 	if err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
+		traceeClient.CloseConnection()
+		return
 	}
-	defer TCS.CloseConnection()
+	defer traceeClient.CloseConnection()
 
 	// create stream from client
 	req := &pb.StreamEventsRequest{Policies: args}
-	stream, err := TCS.StreamEvents(cmd.Context(), req)
+	stream, err := traceeClient.StreamEvents(cmd.Context(), req)
 	if err != nil {
 		cmd.PrintErrln("Error calling Stream: ", err)
-	}
-
-	//add check for the output flag
-	//TODO:support only table and json format for now
-	switch formatFlag {
-	case "json":
-		jsonStreamEvents(cmd, args, stream)
-	case "table":
-		tableStreamEvents(cmd, args, stream)
-	case "template": // go template
-		fallthrough
-	default:
-		cmd.PrintErrln("Error: output format not supported")
 		return
 	}
-}
-
-// tableStreamEvents prints events in a table format
-func tableStreamEvents(cmd *cobra.Command, _ []string, stream pb.TraceeService_StreamEventsClient) {
-	// Init table header before streaming starts
-	tbl := formatter.New(formatFlag, outputFlag, cmd)
-	tbl.PrintTableHeaders()
-	// Receive and process streamed responses
-	for {
-		res, err := stream.Recv()
-		if err != nil {
-			// Handle the error that occurs when the server closes the stream
-			if err.Error() == "EOF" {
-				break
-			}
-			cmd.PrintErrln("Error receiving streamed event: ", err)
-		}
-		tbl.PrintTableRow(res.Event)
 
+	//create formatter for output
+	format, err := formatter.New(formatFlag, outputFlag, cmd)
+	if err != nil {
+		cmd.PrintErrln("Error creating formatter: ", err)
+		return
 	}
-}
+	//show events
+	printer.StreamEvents(format, args, stream)
 
-// jsonStreamEvents prints events in json format
-func jsonStreamEvents(cmd *cobra.Command, _ []string, stream pb.TraceeService_StreamEventsClient) {
-	// Receive and process streamed responses
-	for {
-		res, err := stream.Recv()
-		if err != nil {
-			// Handle the error that occurs when the server closes the stream
-			if err.Error() == "EOF" {
-				break
-			}
-			cmd.PrintErrln("Error receiving streamed event: ", err)
-		}
-		// Print each event as a row in json format
-		formatter.PrintJSON(cmd, res.Event, outputFlag)
-	}
 }