test watchloop

Author: pchila

internal/pkg/agent/cmd/watch.go

@@ -8,9 +8,7 @@ import (
 	"context"
 	"fmt"
 	"os"
-	"os/signal"
 	"runtime"
-	"syscall"
 	"time"
 
 	"github.com/spf13/cobra"
@@ -211,52 +209,6 @@ func isWindows() bool {
 	return runtime.GOOS == "windows"
 }
 
-func watch(ctx context.Context, tilGrace time.Duration, errorCheckInterval time.Duration, log *logger.Logger) error {
-	errChan := make(chan error)
-
-	ctx, cancel := context.WithCancel(ctx)
-
-	//cleanup
-	defer func() {
-		cancel()
-		close(errChan)
-	}()
-
-	agentWatcher := upgrade.NewAgentWatcher(errChan, log, errorCheckInterval)
-	go agentWatcher.Run(ctx)
-
-	signals := make(chan os.Signal, 1)
-	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGHUP)
-
-	t := time.NewTimer(tilGrace)
-	defer t.Stop()
-
-WATCHLOOP:
-	for {
-		select {
-		case s := <-signals:
-			log.Infof("received signal: (%d): %v during watch", s, s)
-			if s == syscall.SIGINT || s == syscall.SIGTERM {
-				log.Infof("received signal: (%d): %v. Exiting watch", s, s)
-				return ErrWatchCancelled
-			}
-			continue
-		case <-ctx.Done():
-			break WATCHLOOP
-		// grace period passed, agent is considered stable
-		case <-t.C:
-			log.Info("Grace period passed, not watching")
-			break WATCHLOOP
-		// Agent in degraded state.
-		case err := <-errChan:
-			log.Errorf("Agent Error detected: %s", err.Error())
-			return err
-		}
-	}
-
-	return nil
-}
-
 // gracePeriod returns true if it is within grace period and time until grace period ends.
 // otherwise it returns false and 0
 func gracePeriod(marker *upgrade.UpdateMarker, gracePeriodDuration time.Duration) (bool, time.Duration) {

internal/pkg/agent/cmd/watch_impl.go

@@ -6,6 +6,9 @@ package cmd
 
 import (
 	"context"
+	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/elastic/elastic-agent-libs/logp"
@@ -29,3 +32,50 @@ func (a upgradeInstallationModifier) Cleanup(log *logger.Logger, topDirPath, cur
 func (a upgradeInstallationModifier) Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPath, prevVersionedHome, prevHash string) error {
 	return upgrade.Rollback(ctx, log, c, topDirPath, prevVersionedHome, prevHash)
 }
+
+func watch(ctx context.Context, tilGrace time.Duration, errorCheckInterval time.Duration, log *logger.Logger) error {
+	errChan := make(chan error)
+
+	ctx, cancel := context.WithCancel(ctx)
+
+	//cleanup
+	defer func() {
+		cancel()
+		close(errChan)
+	}()
+
+	agtWatcher := upgrade.NewAgentWatcher(errChan, log, errorCheckInterval)
+	go agtWatcher.Run(ctx)
+
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
+
+	graceTimer := time.NewTimer(tilGrace)
+	defer graceTimer.Stop()
+
+	return watchLoop(ctx, log, signals, errChan, graceTimer.C)
+}
+
+func watchLoop(ctx context.Context, log *logger.Logger, signals <-chan os.Signal, errChan <-chan error, graceTimer <-chan time.Time) error {
+	for {
+		select {
+		case s := <-signals:
+			log.Infof("received signal: (%d): %v during watch", s, s)
+			if s == syscall.SIGINT || s == syscall.SIGTERM {
+				log.Infof("received signal: (%d): %v. Exiting watch", s, s)
+				return ErrWatchCancelled
+			}
+			continue
+		case <-ctx.Done():
+			return nil
+		// grace period passed, agent is considered stable
+		case <-graceTimer:
+			log.Info("Grace period passed, not watching")
+			return nil
+		// Agent in degraded state.
+		case err := <-errChan:
+			log.Errorf("Agent Error detected: %s", err.Error())
+			return err
+		}
+	}
+}