diff --git a/cmd/traceectl/pkg/cmd/test/test.go b/cmd/traceectl/pkg/cmd/test/test.go
deleted file mode 100644
index 8beab7e7181d..000000000000
--- a/cmd/traceectl/pkg/cmd/test/test.go
+++ /dev/null
@@ -1,52 +0,0 @@
-package test
-
-import (
-	"bytes"
-	"testing"
-	"time"
-
-	"github.com/spf13/cobra"
-	"github.com/stretchr/testify/assert"
-	"google.golang.org/grpc"
-
-	"github.com/aquasecurity/tracee/cmd/traceectl/pkg/mock"
-)
-
-type TestCase struct {
-	TestName        string
-	OutputSlice     []string
-	ExpectedPrinter interface{}
-	ExpectedError   error
-}
-
-func runMockServer(t *testing.T) *grpc.Server {
-	mockServer, err := mock.StartMockServer()
-	if err != nil {
-		t.Fatalf("Failed to start mock server: %v", err)
-	}
-	time.Sleep(100 * time.Millisecond)
-	return mockServer
-}
-func TestCommand(t *testing.T, testCase TestCase, rootCmd *cobra.Command) {
-	server := runMockServer(t)
-	defer server.Stop()
-	var buf bytes.Buffer
-	rootCmd.SetOut(&buf)
-	rootCmd.SetErr(&buf)
-	rootCmd.SetArgs(testCase.OutputSlice)
-	err := rootCmd.Execute()
-	output := buf.String()
-	if err != nil && testCase.ExpectedError == nil {
-		t.Errorf("Unexpected error for test %s: %v", testCase.TestName, err)
-		return
-	}
-	if err == nil && testCase.ExpectedError != nil {
-		t.Errorf("Expected error for test %s but got none", testCase.TestName)
-		return
-	}
-	if testCase.ExpectedError != nil {
-		assert.ErrorContains(t, err, testCase.ExpectedError.Error())
-	} else {
-		assert.Contains(t, output, testCase.ExpectedPrinter)
-	}
-}
diff --git a/cmd/traceectl/pkg/mock/diagnostic_server.go b/cmd/traceectl/pkg/mock/diagnostic_server.go
deleted file mode 100644
index 1d4fcff36e58..000000000000
--- a/cmd/traceectl/pkg/mock/diagnostic_server.go
+++ /dev/null
@@ -1,17 +0,0 @@
-package mock
-
-import (
-	"context"
-
-	pb "github.com/aquasecurity/tracee/api/v1beta1"
-)
-
-var (
-	ExpectedMetrics pb.GetMetricsResponse = pb.GetMetricsResponse{EventCount: 1, EventsFiltered: 2, NetCapCount: 3,
-		BPFLogsCount: 4, ErrorCount: 5, LostEvCount: 6,
-		LostWrCount: 7, LostNtCapCount: 8, LostBPFLogsCount: 9}
-)
-
-func (s *MockDiagnosticServer) GetMetrics(ctx context.Context, req *pb.GetMetricsRequest) (*pb.GetMetricsResponse, error) {
-	return &ExpectedMetrics, nil
-}
diff --git a/cmd/traceectl/pkg/mock/event_server.go b/cmd/traceectl/pkg/mock/event_server.go
deleted file mode 100644
index 689c98d8554c..000000000000
--- a/cmd/traceectl/pkg/mock/event_server.go
+++ /dev/null
@@ -1,28 +0,0 @@
-package mock
-
-import (
-	"context"
-
-	pb "github.com/aquasecurity/tracee/api/v1beta1"
-)
-
-func (s *MockServiceServer) GetEventDefinitions(ctx context.Context, req *pb.GetEventDefinitionsRequest) (*pb.GetEventDefinitionsResponse, error) {
-	var eventDefinition []*pb.EventDefinition
-	for i := 0; i < len(req.EventNames); i++ {
-		eventDefinition = append(eventDefinition, &pb.EventDefinition{Name: req.EventNames[i], Id: (int32(i))})
-	}
-	return &pb.GetEventDefinitionsResponse{Definitions: eventDefinition}, nil
-}
-func (s *MockServiceServer) DescribeEvent(ctx context.Context, req *pb.GetEventDefinitionsRequest) (*pb.GetEventDefinitionsResponse, error) {
-	return s.GetEventDefinitions(ctx, req)
-}
-func (s *MockServiceServer) ListEvent(ctx context.Context, req *pb.GetEventDefinitionsRequest) (*pb.GetEventDefinitionsResponse, error) {
-	return s.GetEventDefinitions(ctx, req)
-}
-
-func (s *MockServiceServer) EnableEvent(ctx context.Context, req *pb.EnableEventRequest) (*pb.EnableEventResponse, error) {
-	return &pb.EnableEventResponse{}, nil
-}
-func (s *MockServiceServer) DisableEvent(ctx context.Context, req *pb.DisableEventRequest) (*pb.DisableEventResponse, error) {
-	return &pb.DisableEventResponse{}, nil
-}
diff --git a/cmd/traceectl/pkg/mock/root_server.go b/cmd/traceectl/pkg/mock/root_server.go
deleted file mode 100644
index 2d9ba32021aa..000000000000
--- a/cmd/traceectl/pkg/mock/root_server.go
+++ /dev/null
@@ -1,11 +0,0 @@
-package mock
-
-import (
-	"context"
-
-	pb "github.com/aquasecurity/tracee/api/v1beta1"
-)
-
-func (s *MockServiceServer) GetVersion(ctx context.Context, req *pb.GetVersionRequest) (*pb.GetVersionResponse, error) {
-	return &pb.GetVersionResponse{Version: ExpectedVersion}, nil
-}
diff --git a/cmd/traceectl/pkg/mock/server.go b/cmd/traceectl/pkg/mock/server.go
deleted file mode 100644
index 8425cc804e90..000000000000
--- a/cmd/traceectl/pkg/mock/server.go
+++ /dev/null
@@ -1,63 +0,0 @@
-package mock
-
-import (
-	"fmt"
-	"net"
-	"os"
-
-	"google.golang.org/grpc"
-
-	pb "github.com/aquasecurity/tracee/api/v1beta1"
-	"github.com/aquasecurity/tracee/cmd/traceectl/pkg/client"
-)
-
-var (
-	ExpectedVersion                    = "v0.22.0-15-gd09d7fca0d"
-	serverInfo      *client.ServerInfo = &client.ServerInfo{
-		Addr:           client.Socket,
-		ConnectionType: client.Protocol_UNIX,
-	}
-)
-
-type MockServiceServer struct {
-	pb.UnimplementedTraceeServiceServer
-}
-type MockDiagnosticServer struct {
-	pb.UnimplementedDiagnosticServiceServer
-}
-
-func CreateMockServer() (*grpc.Server, net.Listener, error) {
-	if _, err := os.Stat(serverInfo.Addr); err == nil {
-		if err := os.Remove(serverInfo.Addr); err != nil {
-			return nil, nil, fmt.Errorf("failed to cleanup gRPC listening address (%s): %v", serverInfo.Addr, err)
-		}
-	}
-	listener, err := net.Listen("unix", serverInfo.Addr)
-	if err != nil {
-		return nil, nil, fmt.Errorf("failed to create Unix socket listener: %v", err)
-	}
-	server := grpc.NewServer()
-	pb.RegisterTraceeServiceServer(server, &MockServiceServer{})
-	pb.RegisterDiagnosticServiceServer(server, &MockDiagnosticServer{})
-
-	return server, listener, nil
-}
-func StartMockServer() (*grpc.Server, error) {
-	mockServer, listener, err := CreateMockServer()
-	if err != nil {
-		return nil, fmt.Errorf("failed to create mock server: %v", err)
-	}
-	go func() {
-		if err := mockServer.Serve(listener); err != nil {
-			fmt.Printf("gRPC server failed: %v\n", err)
-		}
-	}()
-
-	return mockServer, nil
-}
-func StopMockServer(server *grpc.Server) {
-	server.GracefulStop()
-	if err := os.Remove(serverInfo.Addr); err != nil {
-		fmt.Printf("failed to remove Unix socket: %v\n", err)
-	}
-}
diff --git a/cmd/traceectl/pkg/mock/stream_server.go b/cmd/traceectl/pkg/mock/stream_server.go
deleted file mode 100644
index e69db107da21..000000000000
--- a/cmd/traceectl/pkg/mock/stream_server.go
+++ /dev/null
@@ -1,64 +0,0 @@
-package mock
-
-import (
-	"sort"
-	"strings"
-	"time"
-
-	pb "github.com/aquasecurity/tracee/api/v1beta1"
-)
-
-func (s *MockServiceServer) StreamEvents(req *pb.StreamEventsRequest, stream pb.TraceeService_StreamEventsServer) error {
-	mockEvents := CreateEventsFromPolicies(req.Policies)
-	for _, event := range mockEvents {
-		if err := stream.Send(event); err != nil {
-			return err
-		}
-	}
-	time.Sleep(100 * time.Millisecond)
-	return nil
-}
-func generateEvent(policy []string) *pb.Event {
-	return &pb.Event{
-		Policies: &pb.Policies{Matched: policy},
-	}
-}
-func CreateEventsFromPolicies(policies []string) []*pb.StreamEventsResponse {
-	if len(policies) == 0 {
-		return []*pb.StreamEventsResponse{
-			{Event: generateEvent([]string{""})},
-		}
-	}
-	sort.Strings(policies)
-	var results []*pb.StreamEventsResponse
-	combinations := generateCombinations(policies)
-	sort.SliceStable(combinations, func(i, j int) bool {
-		if len(combinations[i]) != len(combinations[j]) {
-			return len(combinations[i]) < len(combinations[j])
-		}
-		return strings.Join(combinations[i], ",") < strings.Join(combinations[j], ",")
-	})
-	for _, combo := range combinations {
-		results = append(results, &pb.StreamEventsResponse{
-			Event: generateEvent(combo),
-		})
-	}
-
-	return results
-}
-func generateCombinations(policies []string) [][]string {
-	var result [][]string
-	n := len(policies)
-	var helper func(start int, combo []string)
-	helper = func(start int, combo []string) {
-		if len(combo) > 0 {
-			combinationCopy := append([]string{}, combo...)
-			result = append(result, combinationCopy)
-		}
-		for i := start; i < n; i++ {
-			helper(i+1, append(combo, policies[i]))
-		}
-	}
-	helper(0, []string{})
-	return result
-}
diff --git a/cmd/traceectl/test/server.go b/cmd/traceectl/test/server.go
new file mode 100644
index 000000000000..7f9bddf60ca1
--- /dev/null
+++ b/cmd/traceectl/test/server.go
@@ -0,0 +1,41 @@
+package test
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"os"
+)
+
+const DefaultSocket = "/tmp/tracee.sock"
+
+type Server struct {
+	addr     string
+	listener net.Listener
+}
+
+func SetupMockSocket() (*Server, error) {
+	listener, err := net.Listen("unix", DefaultSocket)
+	if err != nil {
+		return nil, err
+	}
+	return &Server{
+		addr:     DefaultSocket,
+		listener: listener,
+	}, nil
+}
+
+func (s *Server) TeardownMockSocket() error {
+	if s.listener != nil {
+		if err := s.listener.Close(); err != nil {
+			return fmt.Errorf("failed to close mock socket: %w", err) // Wrap the error
+		}
+	}
+	// Check if the socket file still exists and remove it
+	if _, err := os.Stat(DefaultSocket); err == nil {
+		if err := os.Remove(DefaultSocket); err != nil {
+			log.Printf("Warning: failed to remove mock socket file: %v", err)
+		}
+	}
+	return nil
+}