resolve client creation code

aquasecurity · Dec 8, 2024 · 2159616 · 2159616
1 parent 30317ed
commit 2159616
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 28 deletions.
diff --git a/cmd/traceectl/cmd/event.go b/cmd/traceectl/cmd/event.go
@@ -85,8 +85,8 @@ var disableEventCmd = &cobra.Command{
 }
 
 func listEvents(cmd *cobra.Command, args []string) {
-	var traceeClient client.ServiceClient
-	if err := traceeClient.NewServiceClient(server); err != nil {
+	traceeClient, err := client.NewServiceClient(server)
+	if err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}
@@ -118,8 +118,8 @@ func listEvents(cmd *cobra.Command, args []string) {
 }
 
 func getEventDescriptions(cmd *cobra.Command, args []string) {
-	var traceeClient client.ServiceClient
-	if err := traceeClient.NewServiceClient(server); err != nil {
+	traceeClient, err := client.NewServiceClient(server)
+	if err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}
@@ -161,28 +161,28 @@ func prepareDescription(event *pb.EventDefinition) []string {
 
 }
 func enableEvents(cmd *cobra.Command, eventName string) {
-	var traceeClient client.ServiceClient
-	if err := traceeClient.NewServiceClient(server); err != nil {
+	traceeClient, err := client.NewServiceClient(server)
+	if err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}
 	defer traceeClient.CloseConnection()
 
-	_, err := traceeClient.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)
 		return
 	}
 	cmd.Printf("Enabled event: %s\n", eventName)
 }
 func disableEvents(cmd *cobra.Command, eventName string) {
-	var traceeClient client.ServiceClient
-	if err := traceeClient.NewServiceClient(server); err != nil {
+	traceeClient, err := client.NewServiceClient(server)
+	if err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}
 	defer traceeClient.CloseConnection()
-	_, err := traceeClient.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)
 		return

diff --git a/cmd/traceectl/cmd/root.go b/cmd/traceectl/cmd/root.go
@@ -81,8 +81,8 @@ func Execute() {
 }
 
 func displayMetrics(cmd *cobra.Command, _ []string) {
-	var traceeClient client.DiagnosticClient
-	if err := traceeClient.NewDiagnosticClient(server); err != nil {
+	traceeClient, err := client.NewDiagnosticClient(server)
+	if err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}
@@ -107,8 +107,8 @@ func displayMetrics(cmd *cobra.Command, _ []string) {
 }
 
 func displayVersion(cmd *cobra.Command, _ []string) {
-	var traceeClient client.ServiceClient
-	if err := traceeClient.NewServiceClient(server); err != nil {
+	traceeClient, err := client.NewServiceClient(server)
+	if err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}

diff --git a/cmd/traceectl/cmd/stream.go b/cmd/traceectl/cmd/stream.go
@@ -29,8 +29,7 @@ func init() {
 }
 
 func streamEvents(cmd *cobra.Command, args []string) {
-	var traceeClient client.ServiceClient
-	err := traceeClient.NewServiceClient(server)
+	traceeClient, err := client.NewServiceClient(server)
 	if err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return

diff --git a/cmd/traceectl/pkg/client/diagnostic.go b/cmd/traceectl/pkg/client/diagnostic.go
@@ -12,15 +12,15 @@ type DiagnosticClient struct {
 	client pb.DiagnosticServiceClient
 }
 
-func (tc *DiagnosticClient) NewDiagnosticClient(serverInfo ServerInfo) error {
+func NewDiagnosticClient(serverInfo ServerInfo) (*DiagnosticClient, error) {
 	conn, err := connectToServer(serverInfo)
 	if err != nil {
-		return err
+		return nil, err
 	}
-	tc.conn = conn
-	tc.client = pb.NewDiagnosticServiceClient(conn)
-
-	return nil
+	return &DiagnosticClient{
+		conn:   conn,
+		client: pb.NewDiagnosticServiceClient(conn),
+	}, nil
 }
 func (tc *DiagnosticClient) CloseConnection() {
 	if err := tc.conn.Close(); err != nil {

diff --git a/cmd/traceectl/pkg/client/service.go b/cmd/traceectl/pkg/client/service.go
@@ -12,15 +12,15 @@ type ServiceClient struct {
 	client pb.TraceeServiceClient
 }
 
-func (tc *ServiceClient) NewServiceClient(serverInfo ServerInfo) error {
+func NewServiceClient(serverInfo ServerInfo) (*ServiceClient, error) {
 	conn, err := connectToServer(serverInfo)
 	if err != nil {
-		return err
+		return nil, err
 	}
-	tc.conn = conn
-	tc.client = pb.NewTraceeServiceClient(conn)
-
-	return nil
+	return &ServiceClient{
+		conn:   conn,
+		client: pb.NewTraceeServiceClient(conn),
+	}, nil
 }
 func (tc *ServiceClient) CloseConnection() {
 	if err := tc.conn.Close(); err != nil {