enhancement(5235): added unpack error handling tests in upgrade_test

enhancement(5235): updated the unpack tests

Author: Kaan Yalti

internal/pkg/agent/application/upgrade/upgrade_test.go

@@ -5,9 +5,13 @@
 package upgrade
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
 	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -26,6 +30,8 @@ import (
 	"github.com/elastic/elastic-agent-libs/transport/tlscommon"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/paths"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/upgrade/artifact"
+	downloadErrors "github.com/elastic/elastic-agent/internal/pkg/agent/application/upgrade/artifact/download/errors"
+	"github.com/elastic/elastic-agent/internal/pkg/agent/application/upgrade/common"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/upgrade/details"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/errors"
 	"github.com/elastic/elastic-agent/internal/pkg/config"
@@ -38,6 +44,7 @@ import (
 	"github.com/elastic/elastic-agent/pkg/core/logger"
 	"github.com/elastic/elastic-agent/pkg/core/logger/loggertest"
 	agtversion "github.com/elastic/elastic-agent/pkg/version"
+	"github.com/elastic/elastic-agent/testing/mocks/internal_/pkg/agent/application/info"
 	mocks "github.com/elastic/elastic-agent/testing/mocks/pkg/control/v2/client"
 )
 
@@ -1365,3 +1372,131 @@ func createArchive(t *testing.T, archiveName string, archiveFiles []files) (stri
 	}
 	return createTarArchive(t, archiveName, archiveFiles)
 }
+
+func TestUpgradeUnpackErrors(t *testing.T) {
+	log, _ := loggertest.New("test")
+
+	targetVersion := agtversion.NewParsedSemVer(3, 4, 5, "SNAPSHOT", "")
+	targetArtifactName, targetArchiveFiles := buildArchiveFiles(t, archiveFilesWithMoreComponents, targetVersion, "ghijkl")
+
+	// Get the content of the component file to be used in copy function
+	// assertions. It will be used to make sure the error gets triggered in the
+	// copy call in the unpacker and not the downloader.
+	targetArchiveComponentsFileContent := ""
+	for _, file := range targetArchiveFiles {
+		if file.fType == REGULAR {
+			dir := filepath.Dir(file.path)
+			if strings.HasSuffix(dir, "components") {
+				targetArchiveComponentsFileContent = file.content
+				break
+			}
+		}
+	}
+
+	// dir name of the versioned home
+	// used to calculate component paths
+	newVersionedDirName := "elastic-agent-3.4.5-SNAPSHOT-ghijkl"
+
+	mockAgentInfo := info.NewAgent(t)
+	mockAgentInfo.On("Version").Return(targetVersion.String())
+
+	upgradeDetails := details.NewDetails(targetVersion.String(), details.StateRequested, "test")
+
+	mockStdlibFuncs := []common.MockStdLibFuncName{common.CopyFuncName, common.OpenFileFuncName, common.MkdirAllFuncName}
+
+	testCases := map[string]struct {
+		mockStdlibFunc    common.MockStdLibFuncName
+		mockReturnedError error
+		expectedError     error
+	}{}
+
+	for _, mockStdlibFunc := range mockStdlibFuncs {
+		for _, te := range downloadErrors.OS_DiskSpaceErrors {
+			testCases[fmt.Sprintf("unpack_should_return_error_if_unpack_%s_fails: %v", mockStdlibFunc, te)] = struct {
+				mockStdlibFunc    common.MockStdLibFuncName
+				mockReturnedError error
+				expectedError     error
+			}{
+				mockStdlibFunc:    mockStdlibFunc,
+				mockReturnedError: te,
+				expectedError:     downloadErrors.ErrInsufficientDiskSpace,
+			}
+		}
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			tempDir := t.TempDir()
+			paths.SetTop(tempDir)
+			paths.SetDownloads(filepath.Join(tempDir, "downloads"))
+
+			targetArchive, err := createArchive(t, targetArtifactName, targetArchiveFiles)
+			require.NoError(t, err)
+
+			// Used to assert that the mkdirAll error is triggered in the
+			// unpacker and not the downloader
+			newComponentsDir := filepath.Join(paths.Data(), newVersionedDirName, "components")
+			// Used to assert that the openFile error is triggered in the
+			// unpacker and not the downloader
+			newComponentsFile := filepath.Join(newComponentsDir, "comp1")
+
+			t.Logf("Created archive: %s", targetArchive)
+
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				http.ServeFile(w, r, targetArchive)
+			}))
+			t.Cleanup(server.Close)
+
+			config := artifact.Config{
+				TargetDirectory:        paths.Downloads(),
+				SourceURI:              server.URL,
+				RetrySleepInitDuration: 1 * time.Second,
+				HTTPTransportSettings: httpcommon.HTTPTransportSettings{
+					Timeout: 1 * time.Second,
+				},
+			}
+
+			stdlibMocker := common.PrepareStdLibMocks(common.StdLibMocks{
+				CopyMock: func(dst io.Writer, src io.Reader) (int64, error) {
+					// If the content is not the same as the target archive components file content,
+					// write the content to the destination. This is to make
+					// sure that the error is triggered in unpacker and not in downloader
+					buf, err := io.ReadAll(src)
+					require.NoError(t, err)
+					if !bytes.Equal(buf, []byte(targetArchiveComponentsFileContent)) {
+						return io.Copy(dst, bytes.NewReader(buf))
+					}
+
+					return 0, tc.mockReturnedError
+				},
+				OpenFileMock: func(name string, flag int, perm os.FileMode) (*os.File, error) {
+					// Make sure that the openFile error is triggered in the
+					// unpacker and not the downloader
+					if name != newComponentsFile {
+						return os.OpenFile(name, flag, perm)
+					}
+
+					return nil, tc.mockReturnedError
+				},
+				MkdirAllMock: func(path string, perm os.FileMode) error {
+					// Make sure that the mkdirAll error is triggered in the
+					// unpacker and not the downloader
+					if path != newComponentsDir {
+						return os.MkdirAll(path, perm)
+					}
+
+					return tc.mockReturnedError
+				},
+			})
+
+			stdlibMocker(t, tc.mockStdlibFunc)
+
+			upgrader, err := NewUpgrader(log, &config, mockAgentInfo)
+			require.NoError(t, err)
+
+			_, err = upgrader.Upgrade(context.Background(), targetVersion.String(), server.URL, nil, upgradeDetails, true, true)
+			require.ErrorIs(t, err, tc.expectedError, "expected error mismatch")
+			require.Equal(t, upgradeDetails.State, details.StateExtracting)
+		})
+	}
+}