Add context to RunBinary interface (#5384)

Signed-off-by: khanhtc1202 <[email protected]>

Author: khanhtc1202

pkg/app/launcher/cmd/launcher/launcher.go

@@ -306,7 +306,7 @@ func (l *launcher) run(ctx context.Context, input cli.Input) error {
 		}
 
 		// Start new piped process.
-		runningPiped, err = l.launchNewPiped(version, config, workingDir, input.Logger)
+		runningPiped, err = l.launchNewPiped(ctx, version, config, workingDir, input.Logger)
 		if err != nil {
 			input.Logger.Error("LAUNCHER: failed while launching new Piped", zap.Error(err))
 			return err
@@ -404,7 +404,7 @@ func (l *launcher) cleanOldPiped(cmd *lifecycle.Command, workingDir string, logg
 	return nil
 }
 
-func (l *launcher) launchNewPiped(version string, config []byte, workingDir string, logger *zap.Logger) (*lifecycle.Command, error) {
+func (l *launcher) launchNewPiped(ctx context.Context, version string, config []byte, workingDir string, logger *zap.Logger) (*lifecycle.Command, error) {
 	if err := os.MkdirAll(workingDir, 0755); err != nil {
 		return nil, fmt.Errorf("could not create working directory %s (%w)", workingDir, err)
 	}
@@ -436,7 +436,7 @@ func (l *launcher) launchNewPiped(version string, config []byte, workingDir stri
 	args := makePipedArgs(os.Args[2:], configFilePath)
 	logger.Info(fmt.Sprintf("LAUNCHER: start running Piped %s with args %v", version, args))
 
-	return lifecycle.RunBinary(pipedPath, args)
+	return lifecycle.RunBinary(ctx, pipedPath, args)
 }
 
 func (l *launcher) loadConfigData(ctx context.Context) ([]byte, error) {

pkg/lifecycle/binary.go

@@ -70,8 +70,8 @@ func (c *Command) GracefulStop(period time.Duration) error {
 	}
 }
 
-func RunBinary(execPath string, args []string) (*Command, error) {
-	cmd, err := backoff.NewRetry(runBinaryRetryCount, backoff.NewConstant(5*time.Second)).Do(context.Background(), func() (interface{}, error) {
+func RunBinary(ctx context.Context, execPath string, args []string) (*Command, error) {
+	cmd, err := backoff.NewRetry(runBinaryRetryCount, backoff.NewConstant(5*time.Second)).Do(ctx, func() (interface{}, error) {
 		cmd := exec.Command(execPath, args...)
 		cmd.Stdin = os.Stdin
 		cmd.Stdout = os.Stdout

pkg/lifecycle/binary_test.go

@@ -15,12 +15,16 @@
 package lifecycle
 
 import (
+	"context"
+	"net/http"
+	"net/http/httptest"
 	"strconv"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/zap/zaptest"
 )
 
 func TestGracefulStopCommand(t *testing.T) {
@@ -40,7 +44,7 @@ func TestGracefulStopCommand(t *testing.T) {
 
 	for _, tc := range testcases {
 		t.Run(tc.name, func(t *testing.T) {
-			cmd, err := RunBinary("sh", []string{"sleep", "1m"})
+			cmd, err := RunBinary(context.TODO(), "sh", []string{"sleep", "1m"})
 			require.NoError(t, err)
 			require.NotNil(t, cmd)
 
@@ -71,7 +75,7 @@ func TestGracefulStopCommandResult(t *testing.T) {
 
 	for _, tc := range testcases {
 		t.Run(tc.name, func(t *testing.T) {
-			cmd, err := RunBinary("sh", []string{"-c", "exit " + strconv.Itoa(tc.exitCode)})
+			cmd, err := RunBinary(context.TODO(), "sh", []string{"-c", "exit " + strconv.Itoa(tc.exitCode)})
 			require.NoError(t, err)
 			require.NotNil(t, cmd)
 
@@ -81,3 +85,42 @@ func TestGracefulStopCommandResult(t *testing.T) {
 		})
 	}
 }
+
+func TestDownloadBinary(t *testing.T) {
+	server := httpTestServer()
+	defer server.Close()
+
+	logger := zaptest.NewLogger(t)
+	destDir := t.TempDir()
+	destFile := "test-binary"
+
+	t.Run("successful download", func(t *testing.T) {
+		url := server.URL + "/binary"
+		path, err := DownloadBinary(url, destDir, destFile, logger)
+		require.NoError(t, err)
+		assert.FileExists(t, path)
+	})
+
+	t.Run("file already exists", func(t *testing.T) {
+		url := server.URL + "/binary"
+		path, err := DownloadBinary(url, destDir, destFile, logger)
+		require.NoError(t, err)
+		assert.FileExists(t, path)
+
+		// Try downloading again, should not error and file should still exist
+		path, err = DownloadBinary(url, destDir, destFile, logger)
+		require.NoError(t, err)
+		assert.FileExists(t, path)
+	})
+}
+
+func httpTestServer() *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/binary" {
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte("test binary content"))
+		} else {
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+}