fix ha agent logic

DataDog · Nov 27, 2024 · f65d6d3 · f65d6d3
1 parent 060dc7f
commit f65d6d3
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 29 deletions.
diff --git a/comp/haagent/impl/haagent.go b/comp/haagent/impl/haagent.go
@@ -53,8 +53,8 @@ func (h *haAgentImpl) SetLeader(leaderAgentHostname string) {
 // ShouldRunIntegration return true if the agent integrations is should to run.
 // When ha-agent is disabled, the agent behave as standalone agent (non HA) and will always run all integrations.
 func (h *haAgentImpl) ShouldRunIntegration(integrationName string) bool {
-	if h.Enabled() {
-		return validHaIntegrations[integrationName] && h.isLeader.Load()
+	if h.Enabled() && validHaIntegrations[integrationName] {
+		return h.isLeader.Load()
 	}
 	return true
 }

diff --git a/comp/haagent/impl/haagent_test.go b/comp/haagent/impl/haagent_test.go
@@ -179,7 +179,9 @@ func Test_haAgentImpl_ShouldRunIntegration(t *testing.T) {
 		expectShouldRunIntegration map[string]bool
 	}{
 		{
-			name: "should run: for HA integration",
+			name: "ha agent enabled and agent is leader",
+			// should run HA-integrations
+			// should run "non HA integrations"
 			agentConfigs: map[string]interface{}{
 				"hostname":         testAgentHostname,
 				"ha_agent.enabled": true,
@@ -189,12 +191,14 @@ func Test_haAgentImpl_ShouldRunIntegration(t *testing.T) {
 			expectShouldRunIntegration: map[string]bool{
 				"snmp":                true,
 				"network_path":        true,
-				"unknown_integration": false,
-				"cpu":                 false,
+				"unknown_integration": true,
+				"cpu":                 true,
 			},
 		},
 		{
-			name: "should not run: current agent is not leader",
+			name: "ha agent enabled and agent is not leader",
+			// should skip HA-integrations
+			// should run "non HA integrations"
 			agentConfigs: map[string]interface{}{
 				"hostname":         testAgentHostname,
 				"ha_agent.enabled": true,
@@ -204,12 +208,13 @@ func Test_haAgentImpl_ShouldRunIntegration(t *testing.T) {
 			expectShouldRunIntegration: map[string]bool{
 				"snmp":                false,
 				"network_path":        false,
-				"unknown_integration": false,
-				"cpu":                 false,
+				"unknown_integration": true,
+				"cpu":                 true,
 			},
 		},
 		{
-			name: "should run: when HA Agent not enabled",
+			name: "ha agent not enabled",
+			// should run all integrations
 			agentConfigs: map[string]interface{}{
 				"hostname":         testAgentHostname,
 				"ha_agent.enabled": false,
@@ -230,7 +235,7 @@ func Test_haAgentImpl_ShouldRunIntegration(t *testing.T) {
 			haAgent.Comp.SetLeader(tt.leader)
 
 			for integrationName, shouldRun := range tt.expectShouldRunIntegration {
-				assert.Equal(t, shouldRun, haAgent.Comp.ShouldRunIntegration(integrationName))
+				assert.Equalf(t, shouldRun, haAgent.Comp.ShouldRunIntegration(integrationName), "fail for integration: "+integrationName)
 			}
 		})
 	}

diff --git a/pkg/collector/worker/worker_test.go b/pkg/collector/worker/worker_test.go
@@ -642,29 +642,38 @@ func TestWorker_HaIntegration(t *testing.T) {
 	testHostname := "myhost"
 
 	tests := []struct {
-		name             string
-		haAgentEnabled   bool
-		setLeaderValue   string
-		expectedRunCount int
+		name                         string
+		haAgentEnabled               bool
+		setLeaderValue               string
+		expectedSnmpCheckRunCount    int
+		expectedUnknownCheckRunCount int
 	}{
 		{
-			name:             "should run: ha-agent enabled and is leader",
-			haAgentEnabled:   true,
-			setLeaderValue:   testHostname,
-			expectedRunCount: 1,
+			name: "ha-agent enabled and is leader",
+			// should run HA-integrations
+			// should run "non HA integrations"
+			haAgentEnabled:               true,
+			setLeaderValue:               testHostname,
+			expectedSnmpCheckRunCount:    1,
+			expectedUnknownCheckRunCount: 1,
 		},
 		{
-			name:             "should not run: ha-agent enabled and not leader",
-			haAgentEnabled:   true,
-			setLeaderValue:   "leader-is-another-agent",
-			expectedRunCount: 0,
+			name: "ha-agent enabled and not leader",
+			// should skip HA-integrations
+			// should run "non HA integrations"
+			haAgentEnabled:               true,
+			setLeaderValue:               "leader-is-another-agent",
+			expectedSnmpCheckRunCount:    0,
+			expectedUnknownCheckRunCount: 1,
 		},
 		{
+			name: "ha-agent disabled",
 			// When ha-agent is disabled, the agent behave as standalone agent (non HA) and will always run all integrations.
-			name:             "should run: ha-agent disabled",
-			haAgentEnabled:   false,
-			setLeaderValue:   "",
-			expectedRunCount: 1,
+			// should run all integrations
+			haAgentEnabled:               false,
+			setLeaderValue:               "",
+			expectedSnmpCheckRunCount:    1,
+			expectedUnknownCheckRunCount: 1,
 		},
 	}
 	for _, tt := range tests {
@@ -677,9 +686,11 @@ func TestWorker_HaIntegration(t *testing.T) {
 			pendingChecksChan := make(chan check.Check, 10)
 			mockShouldAddStatsFunc := func(checkid.ID) bool { return true }
 
-			testCheck1 := newCheck(t, "snmp:123", false, nil)
+			snmpCheck := newCheck(t, "snmp:123", false, nil)
+			unknownCheck := newCheck(t, "unknown-check:123", false, nil)
 
-			pendingChecksChan <- testCheck1
+			pendingChecksChan <- snmpCheck
+			pendingChecksChan <- unknownCheck
 			close(pendingChecksChan)
 
 			agentConfigs := map[string]interface{}{
@@ -710,7 +721,8 @@ func TestWorker_HaIntegration(t *testing.T) {
 
 			wg.Wait()
 
-			assert.Equal(t, tt.expectedRunCount, testCheck1.RunCount())
+			assert.Equal(t, tt.expectedSnmpCheckRunCount, snmpCheck.RunCount())
+			assert.Equal(t, tt.expectedUnknownCheckRunCount, unknownCheck.RunCount())
 
 			// make sure the check is deleted from checksTracker
 			assert.Equal(t, 0, len(checksTracker.RunningChecks()))