feature(4890): updated the function signature of the upgrade command,…

… updated tests, added new tests

internal/pkg/agent/cmd/upgrade.go

@@ -42,6 +42,7 @@ func newUpgradeCommandWithArgs(_ []string, streams *cli.IOStreams) *cobra.Comman
 		Long:  "This command upgrades the currently installed Elastic Agent to the specified version.",
 		Args:  cobra.ExactArgs(1),
 		Run: func(c *cobra.Command, args []string) {
+			c.SetContext(context.Background())
 			if err := upgradeCmd(streams, c, args); err != nil {
 				fmt.Fprintf(streams.Err, "Error: %v\n%s\n", err, troubleshootMessage())
 				os.Exit(1)
@@ -61,27 +62,40 @@ func newUpgradeCommandWithArgs(_ []string, streams *cli.IOStreams) *cobra.Comman
 	return cmd
 }
 
-func upgradeCmd(streams *cli.IOStreams, cmd *cobra.Command, args []string) error {
-	c := client.New()
-	return upgradeCmdWithClient(streams, cmd, args, c)
+type upgradeInput struct {
+	streams   *cli.IOStreams
+	cmd       *cobra.Command
+	args      []string
+	c         client.Client
+	agentInfo info.Agent
+	cFunc     confirmFunc
 }
 
-func shouldUpgrade(ctx context.Context, cmd *cobra.Command) (bool, error) {
-	agentInfo, err := info.NewAgentInfoWithLog(ctx, "error", false)
+type confirmFunc func(string, bool) (bool, error)
+
+func upgradeCmd(streams *cli.IOStreams, cmd *cobra.Command, args []string) error {
+	c := client.New()
+	agentInfo, err := info.NewAgentInfoWithLog(cmd.Context(), "error", false)
 	if err != nil {
-		return false, fmt.Errorf("failed to retrieve agent info while tring to upgrade the agent: %w", err)
+		return fmt.Errorf("failed to retrieve agent info while tring to upgrade the agent: %w", err)
+	}
+	input := &upgradeInput{
+		streams,
+		cmd,
+		args,
+		c,
+		agentInfo,
+		cli.Confirm,
 	}
+	return upgradeCmdWithClient(input)
+}
 
+func shouldUpgrade(cmd *cobra.Command, agentInfo info.Agent, cFunc confirmFunc) (bool, error) {
 	if agentInfo.IsStandalone() {
 		return true, nil
 	}
 
-	isAdmin, err := utils.HasRoot()
-	if err != nil {
-		return false, fmt.Errorf("failed checking root/Administrator rights while trying to upgrade the agent: %w", err)
-	}
-
-	if !isAdmin {
+	if agentInfo.Unprivileged() {
 		return false, fmt.Errorf("upgrade command needs to be executed as root for fleet managed agents")
 	}
 
@@ -94,7 +108,7 @@ func shouldUpgrade(ctx context.Context, cmd *cobra.Command) (bool, error) {
 		return false, fmt.Errorf("upgrading fleet managed agents is not supported")
 	}
 
-	cf, err := cli.Confirm("Upgrading fleet managed agents is not supported. Would you still like to proceed?", false)
+	cf, err := cFunc("Upgrading fleet managed agents is not supported. Would you still like to proceed?", false)
 	if err != nil {
 		return false, fmt.Errorf("failed while confirming action: %w", err)
 	}
@@ -106,18 +120,18 @@ func shouldUpgrade(ctx context.Context, cmd *cobra.Command) (bool, error) {
 	return true, nil
 }
 
-func upgradeCmdWithClient(streams *cli.IOStreams, cmd *cobra.Command, args []string, c client.Client) error {
-	version := args[0]
+func upgradeCmdWithClient(input *upgradeInput) error {
+	cmd := input.cmd
+	c := input.c
+	version := input.args[0]
 	sourceURI, _ := cmd.Flags().GetString(flagSourceURI)
 
-	ctx := context.Background()
-
-	su, err := shouldUpgrade(ctx, cmd)
+	su, err := shouldUpgrade(cmd, input.agentInfo, input.cFunc)
 	if !su {
 		return fmt.Errorf("aborting upgrade: %w", err)
 	}
 
-	err = c.Connect(ctx)
+	err = c.Connect(cmd.Context())
 	if err != nil {
 		return errors.New(err, "Failed communicating to running daemon", errors.TypeNetwork, errors.M("socket", control.Address()))
 	}
@@ -173,6 +187,6 @@ func upgradeCmdWithClient(streams *cli.IOStreams, cmd *cobra.Command, args []str
 			return errors.New(err, "Failed trigger upgrade of daemon")
 		}
 	}
-	fmt.Fprintf(streams.Out, "Upgrade triggered to version %s, Elastic Agent is currently restarting\n", version)
+	fmt.Fprintf(input.streams.Out, "Upgrade triggered to version %s, Elastic Agent is currently restarting\n", version)
 	return nil
 }

internal/pkg/agent/cmd/upgrade_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/elastic/elastic-agent/internal/pkg/cli"
 	"github.com/elastic/elastic-agent/pkg/control/v2/client"
 	"github.com/elastic/elastic-agent/pkg/control/v2/cproto"
+	mockinfo "github.com/elastic/elastic-agent/testing/mocks/internal_/pkg/agent/application/info"
 )
 
 func TestUpgradeCmd(t *testing.T) {
@@ -31,6 +32,8 @@ func TestUpgradeCmd(t *testing.T) {
 
 		upgradeCh := make(chan struct{})
 		mock := &mockServer{upgradeStop: upgradeCh}
+		mockAgentInfo := mockinfo.NewAgent(t)
+		mockAgentInfo.On("IsStandalone").Return(true)
 		cproto.RegisterElasticAgentControlServer(s, mock)
 		go func() {
 			err := s.Serve(tcpServer)
@@ -43,10 +46,20 @@ func TestUpgradeCmd(t *testing.T) {
 		args := []string{"--skip-verify", "8.13.0"}
 		streams := cli.NewIOStreams()
 		cmd := newUpgradeCommandWithArgs(args, streams)
+		cmd.SetContext(context.Background())
+
+		commandInput := &upgradeInput{
+			streams,
+			cmd,
+			args,
+			c,
+			mockAgentInfo,
+			nil,
+		}
 
 		// the upgrade command will hang until the server shut down
 		go func() {
-			err = upgradeCmdWithClient(streams, cmd, args, c)
+			err = upgradeCmdWithClient(commandInput)
 			assert.NoError(t, err)
 			// verify that we actually talked to the server
 			counter := atomic.LoadInt32(&mock.upgrades)
@@ -68,6 +81,239 @@ func TestUpgradeCmd(t *testing.T) {
 		// this makes sure all client assertions are done
 		<-clientCh
 	})
+	t.Run("fail if fleet managed and unprivileged", func(t *testing.T) {
+		// Set up mock TCP server for gRPC connection
+		tcpServer, err := net.Listen("tcp", "127.0.0.1:")
+		require.NoError(t, err)
+		defer tcpServer.Close()
+
+		s := grpc.NewServer()
+		defer s.Stop()
+
+		// Define mock server and agent information
+		upgradeCh := make(chan struct{})
+		mock := &mockServer{upgradeStop: upgradeCh}
+		mockAgentInfo := mockinfo.NewAgent(t)
+		mockAgentInfo.On("IsStandalone").Return(false) // Simulate fleet-managed agent
+		mockAgentInfo.On("Unprivileged").Return(true)  // Simulate unprivileged mode
+		cproto.RegisterElasticAgentControlServer(s, mock)
+
+		go func() {
+			err := s.Serve(tcpServer)
+			assert.NoError(t, err)
+		}()
+
+		// Create client and command
+		c := client.New(client.WithAddress("http://" + tcpServer.Addr().String()))
+		args := []string{"8.13.0"} // Version argument
+		streams := cli.NewIOStreams()
+		cmd := newUpgradeCommandWithArgs(args, streams)
+		cmd.SetContext(context.Background())
+
+		commandInput := &upgradeInput{
+			streams,
+			cmd,
+			args,
+			c,
+			mockAgentInfo,
+			nil,
+		}
+
+		clientCh := make(chan struct{})
+
+		// Execute upgrade command and validate shouldUpgrade error
+		go func() {
+			err = upgradeCmdWithClient(commandInput)
+
+			// Expect an error due to unprivileged fleet-managed mode
+			assert.Error(t, err)
+			assert.Contains(t, err.Error(), "upgrade command needs to be executed as root for fleet managed agents")
+
+			// Verify counter has not incremented since upgrade should not proceed
+			counter := atomic.LoadInt32(&mock.upgrades)
+			assert.Equal(t, int32(0), counter, "server should not have handled any upgrades")
+
+			close(clientCh)
+		}()
+
+		<-clientCh // Ensure goroutine completes before ending test
+	})
+
+	t.Run("fail if fleet managed privileged but no force flag", func(t *testing.T) {
+		// Set up mock TCP server for gRPC connection
+		tcpServer, err := net.Listen("tcp", "127.0.0.1:")
+		require.NoError(t, err)
+		defer tcpServer.Close()
+
+		s := grpc.NewServer()
+		defer s.Stop()
+
+		// Define mock server and agent information
+		mock := &mockServer{}
+		mockAgentInfo := mockinfo.NewAgent(t)
+		mockAgentInfo.On("IsStandalone").Return(false) // Simulate fleet-managed agent
+		mockAgentInfo.On("Unprivileged").Return(false) // Simulate privileged mode
+		cproto.RegisterElasticAgentControlServer(s, mock)
+
+		go func() {
+			err := s.Serve(tcpServer)
+			assert.NoError(t, err)
+		}()
+
+		// Create client and command
+		c := client.New(client.WithAddress("http://" + tcpServer.Addr().String()))
+		args := []string{"8.13.0"} // Version argument
+		streams := cli.NewIOStreams()
+		cmd := newUpgradeCommandWithArgs(args, streams)
+		cmd.SetContext(context.Background())
+
+		commandInput := &upgradeInput{
+			streams,
+			cmd,
+			args,
+			c,
+			mockAgentInfo,
+			nil,
+		}
+
+		clientCh := make(chan struct{})
+
+		// Execute upgrade command and validate shouldUpgrade error
+		go func() {
+			err = upgradeCmdWithClient(commandInput)
+
+			// Expect an error due to unprivileged fleet-managed mode
+			assert.Error(t, err)
+			assert.Contains(t, err.Error(), "upgrading fleet managed agents is not supported")
+
+			// Verify counter has not incremented since upgrade should not proceed
+			counter := atomic.LoadInt32(&mock.upgrades)
+			assert.Equal(t, int32(0), counter, "server should not have handled any upgrades")
+
+			close(clientCh)
+		}()
+
+		<-clientCh // Ensure goroutine completes before ending test
+	})
+	t.Run("abort upgrade if fleet managed, privileged, --force is set, and user does not confirm", func(t *testing.T) {
+		// Set up mock TCP server for gRPC connection
+		tcpServer, err := net.Listen("tcp", "127.0.0.1:")
+		require.NoError(t, err)
+		defer tcpServer.Close()
+
+		s := grpc.NewServer()
+		defer s.Stop()
+
+		// Define mock server and agent information
+		mock := &mockServer{}
+		mockAgentInfo := mockinfo.NewAgent(t)
+		mockAgentInfo.On("IsStandalone").Return(false) // Simulate fleet-managed agent
+		mockAgentInfo.On("Unprivileged").Return(false) // Simulate privileged mode
+		cproto.RegisterElasticAgentControlServer(s, mock)
+
+		go func() {
+			err := s.Serve(tcpServer)
+			assert.NoError(t, err)
+		}()
+
+		// Create client and command
+		c := client.New(client.WithAddress("http://" + tcpServer.Addr().String()))
+		args := []string{"8.13.0"} // Version argument
+		streams := cli.NewIOStreams()
+		cmd := newUpgradeCommandWithArgs(args, streams)
+		cmd.SetContext(context.Background())
+		cmd.Flags().Set("force", "true")
+
+		commandInput := &upgradeInput{
+			streams,
+			cmd,
+			args,
+			c,
+			mockAgentInfo,
+			func(s string, b bool) (bool, error) {
+				return false, nil
+			},
+		}
+
+		clientCh := make(chan struct{})
+
+		// Execute upgrade command and validate shouldUpgrade error
+		go func() {
+			err = upgradeCmdWithClient(commandInput)
+
+			// Expect an error because user does not confirm the upgrade
+			assert.Error(t, err)
+			assert.Contains(t, err.Error(), "upgrade not confirmed")
+
+			// Verify counter has not incremented since upgrade should not proceed
+			counter := atomic.LoadInt32(&mock.upgrades)
+			assert.Equal(t, int32(0), counter, "server should not have handled any upgrades")
+
+			close(clientCh)
+		}()
+
+		<-clientCh // Ensure goroutine completes before ending test
+	})
+	t.Run("proceed with upgrade if fleet managed, privileged, --force is set, and user confirms upgrade", func(t *testing.T) {
+		// Set up mock TCP server for gRPC connection
+		tcpServer, err := net.Listen("tcp", "127.0.0.1:")
+		require.NoError(t, err)
+		defer tcpServer.Close()
+
+		s := grpc.NewServer()
+		defer s.Stop()
+
+		// Define mock server and agent information
+		upgradeCh := make(chan struct{})
+		mock := &mockServer{upgradeStop: upgradeCh}
+		mockAgentInfo := mockinfo.NewAgent(t)
+		mockAgentInfo.On("IsStandalone").Return(false) // Simulate fleet-managed agent
+		mockAgentInfo.On("Unprivileged").Return(false) // Simulate privileged mode
+		cproto.RegisterElasticAgentControlServer(s, mock)
+
+		go func() {
+			err := s.Serve(tcpServer)
+			assert.NoError(t, err)
+		}()
+
+		// Create client and command
+		c := client.New(client.WithAddress("http://" + tcpServer.Addr().String()))
+		args := []string{"8.13.0"} // Version argument
+		streams := cli.NewIOStreams()
+		cmd := newUpgradeCommandWithArgs(args, streams)
+		cmd.SetContext(context.Background())
+		cmd.Flags().Set("force", "true")
+
+		commandInput := &upgradeInput{
+			streams,
+			cmd,
+			args,
+			c,
+			mockAgentInfo,
+			func(s string, b bool) (bool, error) {
+				return true, nil
+			},
+		}
+
+		clientCh := make(chan struct{})
+
+		// Execute upgrade command and validate shouldUpgrade error
+		go func() {
+			err = upgradeCmdWithClient(commandInput)
+
+			assert.NoError(t, err)
+
+			// Verify counter has not incremented since upgrade should not proceed
+			counter := atomic.LoadInt32(&mock.upgrades)
+			assert.Equal(t, int32(1), counter, "server should handle exactly one upgrade")
+
+			close(clientCh)
+		}()
+
+		close(upgradeCh)
+
+		<-clientCh // Ensure goroutine completes before ending test
+	})
 }
 
 type mockServer struct {