Refine agent run validation (#513)

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

Author: gkeesh7

agent/cmd/cmd.go

@@ -108,15 +108,7 @@ func WithEffect(f func()) Option {
 
 // Run runs the agent.
 func Run(flags *Flags, opts ...Option) {
-	if flags.PeerPort == 0 {
-		panic("must specify non-zero peer port")
-	}
-	if flags.AgentServerPort == 0 {
-		panic("must specify non-zero agent server port")
-	}
-	if flags.AgentRegistryPort == 0 {
-		panic("must specify non-zero agent registry port")
-	}
+	validateRequiredPorts(flags)
 
 	var overrides options
 	for _, o := range opts {
@@ -272,6 +264,19 @@ func Run(flags *Flags, opts ...Option) {
 	}
 }
 
+// validateRequiredPorts panics if any required port flags are not set.
+func validateRequiredPorts(flags *Flags) {
+	if flags.PeerPort == 0 {
+		panic("must specify non-zero peer port")
+	}
+	if flags.AgentServerPort == 0 {
+		panic("must specify non-zero agent server port")
+	}
+	if flags.AgentRegistryPort == 0 {
+		panic("must specify non-zero agent registry port")
+	}
+}
+
 // heartbeatTicker provides the minimal ticker contract required by heartbeat.
 type heartbeatTicker interface {
 	Chan() <-chan time.Time

agent/cmd/cmd_test.go

@@ -113,6 +113,56 @@ func TestRunValidation(t *testing.T) {
 		})
 	}
 }
+
+func TestValidateRequiredPorts(t *testing.T) {
+	tests := []struct {
+		desc        string
+		flags       Flags
+		shouldPanic bool
+		panicMsg    string
+	}{
+		{
+			desc: "valid ports",
+			flags: Flags{
+				PeerPort:          1,
+				AgentServerPort:   2,
+				AgentRegistryPort: 3,
+			},
+		},
+		{
+			desc:        "missing peer port",
+			flags:       Flags{AgentServerPort: 1, AgentRegistryPort: 1},
+			shouldPanic: true,
+			panicMsg:    "must specify non-zero peer port",
+		},
+		{
+			desc:        "missing agent server port",
+			flags:       Flags{PeerPort: 1, AgentRegistryPort: 1},
+			shouldPanic: true,
+			panicMsg:    "must specify non-zero agent server port",
+		},
+		{
+			desc:        "missing agent registry port",
+			flags:       Flags{PeerPort: 1, AgentServerPort: 1},
+			shouldPanic: true,
+			panicMsg:    "must specify non-zero agent registry port",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			if test.shouldPanic {
+				assert.PanicsWithValue(t, test.panicMsg, func() {
+					validateRequiredPorts(&test.flags)
+				})
+			} else {
+				assert.NotPanics(t, func() {
+					validateRequiredPorts(&test.flags)
+				})
+			}
+		})
+	}
+}
 func TestHeartbeatWithTicker(t *testing.T) {
 	scope := tally.NewTestScope("", nil)
 	mockClock := clock.NewMock()