From dd19b4312b1b704be53ed476e347f39e34062d87 Mon Sep 17 00:00:00 2001
From: ShohamBit <shohambit@gmail.com>
Date: Sun, 8 Dec 2024 11:42:25 +0000
Subject: [PATCH] resolve client creation code

---
 cmd/traceectl/cmd/event.go          | 20 ++++++++++----------
 cmd/traceectl/cmd/root.go           |  4 ++--
 cmd/traceectl/cmd/stream.go         |  3 +--
 cmd/traceectl/pkg/client/service.go | 12 ++++++------
 4 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/cmd/traceectl/cmd/event.go b/cmd/traceectl/cmd/event.go
index 126fd56c90d9..c4abe945d913 100644
--- 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,14 +161,14 @@ 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
@@ -176,13 +176,13 @@ func enableEvents(cmd *cobra.Command, eventName string) {
 	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
index 07edf0de9d26..9e818022c1ef 100644
--- a/cmd/traceectl/cmd/root.go
+++ b/cmd/traceectl/cmd/root.go
@@ -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
index 7a0d25383361..0b9f8b3fd42d 100644
--- 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/service.go b/cmd/traceectl/pkg/client/service.go
index 900f23ce5918..0e6c46a184d6 100644
--- 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 {