Adding tests to cmd.go and refactoring the heartbeat method (AI generated) (#508)

* Added unit tests and refactored cmd.go for kraken agent

* Removed PR template

* Refactor heartbeat ticker for testing

* Split option tests

* Inject heartbeat ticker from Run

* Calling stopheartbeat immediately after hearbeat with a defer

---------

Co-authored-by: gauravk <gauravk@uber.com>

Author: gkeesh7

agent/cmd/cmd.go

@@ -17,6 +17,7 @@ import (
 	"flag"
 	"fmt"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/uber/kraken/agent/agentserver"
@@ -231,32 +232,73 @@ func Run(flags *Flags, opts ...Option) {
 		config.AgentServer, stats, cads, sched, tagClient, announceClient, containerRuntimeFactory)
 	addr := fmt.Sprintf(":%d", flags.AgentServerPort)
 	log.Infof("Starting agent server on %s", addr)
+	heartbeatTicker := &timeTicker{inner: time.NewTicker(10 * time.Second)}
+	heartbeatDone := make(chan struct{})
+	var heartbeatStop sync.Once
+	stopHeartbeat := func() {
+		heartbeatStop.Do(func() {
+			close(heartbeatDone)
+			heartbeatTicker.Stop()
+		})
+	}
+
+	go heartbeat(stats, heartbeatTicker, heartbeatDone)
+	defer stopHeartbeat()
 	go func() {
-		log.Fatal(http.ListenAndServe(addr, agentServer.Handler()))
+		if err := http.ListenAndServe(addr, agentServer.Handler()); err != nil {
+			stopHeartbeat()
+			log.Fatal(err)
+		}
 	}()
 
 	log.Info("Starting registry...")
 	go func() {
-		log.Fatal(registry.ListenAndServe())
+		if err := registry.ListenAndServe(); err != nil {
+			stopHeartbeat()
+			log.Fatal(err)
+		}
 	}()
 
-	go heartbeat(stats)
-
-	log.Fatal(nginx.Run(config.Nginx, map[string]interface{}{
+	if err := nginx.Run(config.Nginx, map[string]interface{}{
 		"allowed_cidrs": config.AllowedCidrs,
 		"port":          flags.AgentRegistryPort,
 		"registry_server": nginx.GetServer(
 			config.Registry.Docker.HTTP.Net, config.Registry.Docker.HTTP.Addr),
 		"agent_server":    fmt.Sprintf("127.0.0.1:%d", flags.AgentServerPort),
 		"registry_backup": config.RegistryBackup},
-		nginx.WithTLS(config.TLS)))
+		nginx.WithTLS(config.TLS)); err != nil {
+		stopHeartbeat()
+		log.Fatal(err)
+	}
+}
+
+// heartbeatTicker provides the minimal ticker contract required by heartbeat.
+type heartbeatTicker interface {
+	Chan() <-chan time.Time
+	Stop()
+}
+
+type timeTicker struct {
+	inner *time.Ticker
+}
+
+func (t *timeTicker) Chan() <-chan time.Time {
+	return t.inner.C
+}
+
+func (t *timeTicker) Stop() {
+	t.inner.Stop()
 }
 
 // heartbeat periodically emits a counter metric which allows us to monitor the
-// number of active agents.
-func heartbeat(stats tally.Scope) {
+// number of active agents, using the provided ticker and done channel to control its lifecycle.
+func heartbeat(stats tally.Scope, ticker heartbeatTicker, done <-chan struct{}) {
 	for {
-		stats.Counter("heartbeat").Inc(1)
-		time.Sleep(10 * time.Second)
+		select {
+		case <-ticker.Chan():
+			stats.Counter("heartbeat").Inc(1)
+		case <-done:
+			return
+		}
 	}
 }

agent/cmd/cmd_test.go

@@ -0,0 +1,158 @@
+package cmd
+
+import (
+	"flag"
+	"os"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/andres-erbsen/clock"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/uber-go/tally"
+	"go.uber.org/zap"
+)
+
+func TestParseFlags(t *testing.T) {
+	// Save original args and flagset
+	oldArgs := os.Args
+	oldCommandLine := flag.CommandLine
+	defer func() {
+		os.Args = oldArgs
+		flag.CommandLine = oldCommandLine
+	}()
+
+	// Reset flags
+	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
+
+	// Set up test args
+	os.Args = []string{
+		"cmd",
+		"-peer-ip=1.2.3.4",
+		"-peer-port=1000",
+		"-agent-server-port=2000",
+		"-agent-registry-port=3000",
+		"-config=config.yaml",
+		"-zone=test-zone",
+		"-cluster=test-cluster",
+		"-secrets=secrets.yaml",
+	}
+
+	flags := ParseFlags()
+
+	assert.Equal(t, "1.2.3.4", flags.PeerIP)
+	assert.Equal(t, 1000, flags.PeerPort)
+	assert.Equal(t, 2000, flags.AgentServerPort)
+	assert.Equal(t, 3000, flags.AgentRegistryPort)
+	assert.Equal(t, "config.yaml", flags.ConfigFile)
+	assert.Equal(t, "test-zone", flags.Zone)
+	assert.Equal(t, "test-cluster", flags.KrakenCluster)
+	assert.Equal(t, "secrets.yaml", flags.SecretsFile)
+}
+
+func TestWithConfigOption(t *testing.T) {
+	var o options
+	c := Config{RegistryBackup: "test"}
+	WithConfig(c)(&o)
+	assert.Equal(t, "test", o.config.RegistryBackup)
+}
+
+func TestWithMetricsOption(t *testing.T) {
+	var o options
+	s := tally.NoopScope
+	WithMetrics(s)(&o)
+	assert.Equal(t, s, o.metrics)
+}
+
+func TestWithLoggerOption(t *testing.T) {
+	var o options
+	l := zap.NewNop()
+	WithLogger(l)(&o)
+	assert.Equal(t, l, o.logger)
+}
+
+func TestWithEffectOption(t *testing.T) {
+	var o options
+	called := false
+	f := func() { called = true }
+	WithEffect(f)(&o)
+	assert.NotNil(t, o.effect)
+	o.effect()
+	assert.True(t, called)
+}
+
+func TestRunValidation(t *testing.T) {
+	tests := []struct {
+		desc  string
+		flags Flags
+		panic string
+	}{
+		{
+			desc:  "missing peer port",
+			flags: Flags{AgentServerPort: 1, AgentRegistryPort: 1},
+			panic: "must specify non-zero peer port",
+		},
+		{
+			desc:  "missing agent server port",
+			flags: Flags{PeerPort: 1, AgentRegistryPort: 1},
+			panic: "must specify non-zero agent server port",
+		},
+		{
+			desc:  "missing agent registry port",
+			flags: Flags{PeerPort: 1, AgentServerPort: 1},
+			panic: "must specify non-zero agent registry port",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			assert.PanicsWithValue(t, test.panic, func() {
+				Run(&test.flags)
+			})
+		})
+	}
+}
+func TestHeartbeatWithTicker(t *testing.T) {
+	scope := tally.NewTestScope("", nil)
+	mockClock := clock.NewMock()
+	mockTicker := mockClock.Ticker(100 * time.Millisecond)
+	done := make(chan struct{})
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		heartbeat(scope, clockTicker{ticker: mockTicker}, done)
+	}()
+
+	for i := 0; i < 3; i++ {
+		mockClock.Add(100 * time.Millisecond)
+	}
+
+	require.Eventually(t, func() bool {
+		snapshot := scope.Snapshot()
+		for _, counter := range snapshot.Counters() {
+			if counter.Name() == "heartbeat" && counter.Value() >= 3 {
+				return true
+			}
+		}
+		return false
+	}, time.Second, 10*time.Millisecond)
+
+	close(done)
+	mockTicker.Stop()
+	wg.Wait()
+}
+
+type clockTicker struct {
+	ticker *clock.Ticker
+}
+
+func (t clockTicker) Chan() <-chan time.Time {
+	return t.ticker.C
+}
+
+func (t clockTicker) Stop() {
+	t.ticker.Stop()
+}