Skip to content

Commit 97ecba8

Browse files
pchilapchila
committed
Add tests for RollbackWithOpts
1 parent 9d0bcd0 commit 97ecba8

File tree

1 file changed

+215
-5
lines changed

1 file changed

+215
-5
lines changed

‎internal/pkg/agent/application/upgrade/rollback_test.go‎

Lines changed: 215 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ package upgrade
66

77
import (
88
"context"
9+
"errors"
910
"fmt"
1011
"os"
1112
"path/filepath"
@@ -16,6 +17,8 @@ import (
1617
"github.com/stretchr/testify/assert"
1718
"github.com/stretchr/testify/mock"
1819
"github.com/stretchr/testify/require"
20+
"go.uber.org/zap/zapcore"
21+
"go.uber.org/zap/zaptest/observer"
1922

2023
"github.com/elastic/elastic-agent/internal/pkg/agent/application/paths"
2124
"github.com/elastic/elastic-agent/pkg/core/logger"
@@ -208,7 +211,7 @@ func TestCleanup(t *testing.T) {
208211
t.Run(name, func(t *testing.T) {
209212
testLogger, _ := loggertest.New(t.Name())
210213
testTop := t.TempDir()
211-
setupAgents(t, testLogger, testTop, tt.agentInstallsSetup)
214+
setupAgents(t, testLogger, testTop, tt.agentInstallsSetup, true)
212215
if tt.additionalSetup != nil {
213216
tt.additionalSetup(t, testTop)
214217
}
@@ -303,7 +306,7 @@ func TestRollback(t *testing.T) {
303306
t.Run(name, func(t *testing.T) {
304307
testLogger, _ := loggertest.New(t.Name())
305308
testTop := t.TempDir()
306-
setupAgents(t, testLogger, testTop, tt.agentInstallsSetup)
309+
setupAgents(t, testLogger, testTop, tt.agentInstallsSetup, true)
307310
if tt.additionalSetup != nil {
308311
tt.additionalSetup(t, testTop)
309312
}
@@ -329,6 +332,210 @@ func TestRollback(t *testing.T) {
329332
}
330333
}
331334

335+
func TestRollbackWithOpts(t *testing.T) {
336+
type hookFuncWithLogs func(t *testing.T, logs *observer.ObservedLogs, topDir string)
337+
338+
agentExecutableName := agentName
339+
if runtime.GOOS == "windows" {
340+
agentExecutableName += ".exe"
341+
}
342+
343+
type args struct {
344+
prevVersionedHome string
345+
prevHash string
346+
rollbackOptions []RollbackOpt
347+
}
348+
349+
tests := map[string]struct {
350+
agentInstallsSetup setupAgentInstallations
351+
setupMocks func(*mocks.Client)
352+
args args
353+
wantErr assert.ErrorAssertionFunc
354+
checkAfterRollback hookFuncWithLogs
355+
}{
356+
"SkipCleanup: leave the current installation intact": {
357+
agentInstallsSetup: setupAgentInstallations{
358+
installedAgents: []testAgentInstall{
359+
{
360+
version: version123Snapshot,
361+
useVersionInPath: true,
362+
},
363+
{
364+
version: version456Snapshot,
365+
useVersionInPath: true,
366+
},
367+
},
368+
},
369+
setupMocks: func(mockClient *mocks.Client) {
370+
mockClient.EXPECT().Connect(
371+
mock.AnythingOfType("*context.timerCtx"),
372+
mock.AnythingOfType("*grpc.funcDialOption"),
373+
mock.AnythingOfType("*grpc.funcDialOption"),
374+
).Return(nil)
375+
mockClient.EXPECT().Disconnect().Return()
376+
mockClient.EXPECT().Restart(mock.Anything).Return(nil).Once()
377+
},
378+
args: args{
379+
prevVersionedHome: "data/elastic-agent-1.2.3-SNAPSHOT-abcdef",
380+
prevHash: "abcdef",
381+
rollbackOptions: []RollbackOpt{
382+
func(rs *RollbackSettings) { rs.SetSkipCleanup(true) },
383+
},
384+
},
385+
wantErr: assert.NoError,
386+
checkAfterRollback: func(t *testing.T, _ *observer.ObservedLogs, topDir string) {
387+
assertAgentInstallExists(t, filepath.Join(topDir, "data", "elastic-agent-1.2.3-SNAPSHOT-abcdef"), agentExecutableName)
388+
assertAgentInstallExists(t, filepath.Join(topDir, "data", "elastic-agent-4.5.6-SNAPSHOT-ghijkl"), agentExecutableName)
389+
linkTarget, err := os.Readlink(filepath.Join(topDir, agentExecutableName))
390+
assert.NoError(t, err, "reading topPath elastic-agent link")
391+
assert.Equal(t, linkTarget, filepath.Join(topDir, "data", "elastic-agent-1.2.3-SNAPSHOT-abcdef", agentExecutableName))
392+
},
393+
},
394+
"SkipRestart: no cleanup, no restart": {
395+
agentInstallsSetup: setupAgentInstallations{
396+
installedAgents: []testAgentInstall{
397+
{
398+
version: version123Snapshot,
399+
useVersionInPath: true,
400+
},
401+
{
402+
version: version456Snapshot,
403+
useVersionInPath: true,
404+
},
405+
},
406+
},
407+
setupMocks: func(mockClient *mocks.Client) {
408+
// nothing to do here, no restart will be issued
409+
},
410+
args: args{
411+
prevVersionedHome: "data/elastic-agent-1.2.3-SNAPSHOT-abcdef",
412+
prevHash: "abcdef",
413+
rollbackOptions: []RollbackOpt{
414+
func(rs *RollbackSettings) { rs.SetSkipRestart(true) },
415+
},
416+
},
417+
wantErr: assert.NoError,
418+
checkAfterRollback: func(t *testing.T, _ *observer.ObservedLogs, topDir string) {
419+
assertAgentInstallExists(t, filepath.Join(topDir, "data", "elastic-agent-1.2.3-SNAPSHOT-abcdef"), agentExecutableName)
420+
assertAgentInstallExists(t, filepath.Join(topDir, "data", "elastic-agent-4.5.6-SNAPSHOT-ghijkl"), agentExecutableName)
421+
linkTarget, err := os.Readlink(filepath.Join(topDir, agentExecutableName))
422+
assert.NoError(t, err, "reading topPath elastic-agent link")
423+
assert.Equal(t, linkTarget, filepath.Join(topDir, "data", "elastic-agent-1.2.3-SNAPSHOT-abcdef", agentExecutableName))
424+
},
425+
},
426+
"Prerestart hook not fatal error: rollback, cleanup and restart as normal": {
427+
agentInstallsSetup: setupAgentInstallations{
428+
installedAgents: []testAgentInstall{
429+
{
430+
version: version123Snapshot,
431+
useVersionInPath: true,
432+
},
433+
{
434+
version: version456Snapshot,
435+
useVersionInPath: true,
436+
},
437+
},
438+
},
439+
setupMocks: func(mockClient *mocks.Client) {
440+
mockClient.EXPECT().Connect(
441+
mock.AnythingOfType("*context.timerCtx"),
442+
mock.AnythingOfType("*grpc.funcDialOption"),
443+
mock.AnythingOfType("*grpc.funcDialOption"),
444+
).Return(nil)
445+
mockClient.EXPECT().Disconnect().Return()
446+
mockClient.EXPECT().Restart(mock.Anything).Return(nil).Once()
447+
},
448+
args: args{
449+
prevVersionedHome: "data/elastic-agent-1.2.3-SNAPSHOT-abcdef",
450+
prevHash: "abcdef",
451+
rollbackOptions: []RollbackOpt{
452+
func(rs *RollbackSettings) {
453+
rs.SetPreRestartHook(func(ctx context.Context, log *logger.Logger, topDirPath string) error {
454+
return errors.New("pre-restart hook error, not fatal")
455+
})
456+
},
457+
},
458+
},
459+
wantErr: assert.NoError,
460+
checkAfterRollback: func(t *testing.T, logs *observer.ObservedLogs, topDir string) {
461+
assertAgentInstallExists(t, filepath.Join(topDir, "data", "elastic-agent-1.2.3-SNAPSHOT-abcdef"), agentExecutableName)
462+
assertAgentInstallCleaned(t, filepath.Join(topDir, "data", "elastic-agent-4.5.6-SNAPSHOT-ghijkl"), agentExecutableName)
463+
linkTarget, err := os.Readlink(filepath.Join(topDir, agentExecutableName))
464+
assert.NoError(t, err, "reading topPath elastic-agent link")
465+
assert.Equal(t, linkTarget, filepath.Join(topDir, "data", "elastic-agent-1.2.3-SNAPSHOT-abcdef", agentExecutableName))
466+
snippetLogs := logs.FilterMessageSnippet("pre-restart hook error, not fatal").All()
467+
if assert.Len(t, snippetLogs, 1) {
468+
assert.Equal(t, zapcore.WarnLevel, snippetLogs[0].Level)
469+
}
470+
},
471+
},
472+
"Prerestart hook fatal error: rollback then return the error (no restart, no cleanup)": {
473+
agentInstallsSetup: setupAgentInstallations{
474+
installedAgents: []testAgentInstall{
475+
{
476+
version: version123Snapshot,
477+
useVersionInPath: true,
478+
},
479+
{
480+
version: version456Snapshot,
481+
useVersionInPath: true,
482+
},
483+
},
484+
},
485+
setupMocks: func(mockClient *mocks.Client) {
486+
// no restart request should be made
487+
},
488+
args: args{
489+
prevVersionedHome: "data/elastic-agent-1.2.3-SNAPSHOT-abcdef",
490+
prevHash: "abcdef",
491+
rollbackOptions: []RollbackOpt{
492+
func(rs *RollbackSettings) {
493+
rs.SetPreRestartHook(func(ctx context.Context, log *logger.Logger, topDirPath string) error {
494+
return fmt.Errorf("fatal pre-restart hook error: %w", FatalRollbackError)
495+
})
496+
},
497+
},
498+
},
499+
wantErr: assert.Error,
500+
checkAfterRollback: func(t *testing.T, logs *observer.ObservedLogs, topDir string) {
501+
assertAgentInstallExists(t, filepath.Join(topDir, "data", "elastic-agent-1.2.3-SNAPSHOT-abcdef"), agentExecutableName)
502+
assertAgentInstallExists(t, filepath.Join(topDir, "data", "elastic-agent-4.5.6-SNAPSHOT-ghijkl"), agentExecutableName)
503+
linkTarget, err := os.Readlink(filepath.Join(topDir, agentExecutableName))
504+
assert.NoError(t, err, "reading topPath elastic-agent link")
505+
assert.Equal(t, linkTarget, filepath.Join(topDir, "data", "elastic-agent-1.2.3-SNAPSHOT-abcdef", agentExecutableName))
506+
},
507+
},
508+
}
509+
for name, tt := range tests {
510+
t.Run(name, func(t *testing.T) {
511+
testLogger, obsLogs := loggertest.New(t.Name())
512+
testTop := t.TempDir()
513+
setupAgents(t, testLogger, testTop, tt.agentInstallsSetup, false)
514+
515+
// mock client
516+
mockClient := mocks.NewClient(t)
517+
tt.setupMocks(mockClient)
518+
519+
tt.wantErr(t, RollbackWithOpts(t.Context(), testLogger, mockClient, testTop, tt.args.prevVersionedHome, tt.args.prevHash, tt.args.rollbackOptions...))
520+
tt.checkAfterRollback(t, obsLogs, testTop)
521+
})
522+
}
523+
}
524+
525+
func assertAgentInstallExists(t *testing.T, versionedHome string, agentExecutableName string) {
526+
assert.DirExists(t, versionedHome)
527+
assert.FileExists(t, filepath.Join(versionedHome, agentExecutableName))
528+
assert.DirExists(t, filepath.Join(versionedHome, "logs"))
529+
assert.DirExists(t, filepath.Join(versionedHome, "run"))
530+
}
531+
532+
func assertAgentInstallCleaned(t *testing.T, versionedHome string, agentExecutableName string) {
533+
assert.DirExists(t, versionedHome)
534+
assert.DirExists(t, filepath.Join(versionedHome, "logs"))
535+
assert.NoDirExists(t, filepath.Join(versionedHome, "run"))
536+
assert.NoFileExists(t, filepath.Join(versionedHome, agentExecutableName))
537+
}
538+
332539
// checkFilesAfterCleanup is a convenience function to check the file structure within topDir.
333540
// *AgentHome paths must be the expected old and new agent paths relative to topDir (typically in the form of "data/elastic-agent-*")
334541
func checkFilesAfterCleanup(t *testing.T, topDir, newAgentHome string, oldAgentHomes ...string) {
@@ -395,7 +602,7 @@ func checkFilesAfterRollback(t *testing.T, topDir, oldAgentHome, newAgentHome st
395602
}
396603

397604
// setupAgents create fake agent installs, update marker file and symlink pointing to one of the installations' elastic-agent placeholder
398-
func setupAgents(t *testing.T, log *logger.Logger, topDir string, installations setupAgentInstallations) {
605+
func setupAgents(t *testing.T, log *logger.Logger, topDir string, installations setupAgentInstallations, writeUpdateMarker bool) {
399606

400607
var (
401608
oldAgentVersion testAgentVersion
@@ -426,8 +633,11 @@ func setupAgents(t *testing.T, log *logger.Logger, topDir string, installations
426633
}
427634
}
428635

429-
t.Logf("Creating upgrade marker from %+v located at %q to %+v located at %q", oldAgentVersion, oldAgentVersionedHome, newAgentVersion, newAgentVersionedHome)
430-
createUpdateMarker(t, log, topDir, newAgentVersion.version, newAgentVersion.hash, newAgentVersionedHome, oldAgentVersion.version, oldAgentVersion.hash, oldAgentVersionedHome, useNewMarker)
636+
if writeUpdateMarker {
637+
t.Logf("Creating upgrade marker from %+v located at %q to %+v located at %q", oldAgentVersion, oldAgentVersionedHome, newAgentVersion, newAgentVersionedHome)
638+
createUpdateMarker(t, log, topDir, newAgentVersion.version, newAgentVersion.hash, newAgentVersionedHome, oldAgentVersion.version, oldAgentVersion.hash, oldAgentVersionedHome, useNewMarker)
639+
}
640+
431641
}
432642

433643
// createFakeAgentInstall will create a mock agent install within topDir, possibly using the version in the directory name, depending on useVersionInPath

0 commit comments

Comments
 (0)