Test CloudFormation agents tags (#1202)

deploy/test-environments/fleet_api/src/agents_enrolled.py

@@ -4,6 +4,8 @@
 """
 import sys
 import time
+import re
+from dataclasses import dataclass
 from api.agent_policy_api import get_agents
 import configuration_fleet as cnfg
 from state_file_manager import state_manager
@@ -12,6 +14,16 @@
 TIMEOUT = 600
 
 
+@dataclass
+class AgentExpected:
+    """
+    Class to represent the details of an enrolled agent.
+    """
+
+    count: int
+    tags: list[str]
+
+
 def get_expected_agents() -> dict:
     """
     Returns:
@@ -20,7 +32,7 @@ def get_expected_agents() -> dict:
     logger.info("Loading agent policies state file")
     policies_dict = {}
     for policy in state_manager.get_policies():
-        policies_dict[policy.agnt_policy_id] = policy.expected_agents
+        policies_dict[policy.agnt_policy_id] = AgentExpected(policy.expected_agents, policy.expected_tags)
     return policies_dict
 
 
@@ -36,23 +48,57 @@ def get_actual_agents() -> dict:
     return policies_dict
 
 
-def verify_agents_enrolled() -> bool:
+def verify_agent_count(expected: dict, actual: dict) -> bool:
     """
     Verify that the expected number of agents are enrolled
     """
-    expected = get_expected_agents()
-    actual = get_actual_agents()
     result = True
-    for policy_id, expected_count in expected.items():
+    for policy_id, expected_agents in expected.items():
         if policy_id not in actual:
             result = False
             logger.info(f"Policy {policy_id} not found in the actual agents mapping")
-        elif actual[policy_id] != expected_count:
+        elif actual[policy_id] != expected_agents.count:
             result = False
-            logger.info(f"Policy {policy_id} expected {expected_count} agents, but got {actual[policy_id]}")
+            logger.info(f"Policy {policy_id} expected {expected_agents.count} agents, but got {actual[policy_id]}")
+        else:
+            logger.info(f"Policy {policy_id} has {actual[policy_id]} agents as expected")
     return result
 
 
+def verify_agent_tags(agent, expected_agents) -> bool:
+    """
+    Verify that the agent has the expected tags
+    """
+    expected_tags = []
+    if agent.policy_id in expected_agents:
+        expected_tags = expected_agents[agent.policy_id].tags
+    for pattern in expected_tags:
+        pattern_exist = False
+        for tag in agent.tags:
+            if re.match(pattern, tag):
+                pattern_exist = True
+                break
+        if not pattern_exist:
+            logger.warning(f"Agent {agent.id} does not have the expected tag {pattern}")
+            return False
+    return True
+
+
+def verify_agents_enrolled() -> bool:
+    """
+    Construct a dictionary of the expected agents and the actual agents
+    Returns:
+        bool: True if the expected agents are enrolled, False otherwise
+    """
+    expected = get_expected_agents()
+    agents = get_agents(cfg=cnfg.elk_config)
+    actual = {}
+    for agent in agents:
+        if verify_agent_tags(agent, expected):
+            actual[agent.policy_id] = actual.get(agent.policy_id, 0) + 1
+    return verify_agent_count(expected, actual)
+
+
 def wait_for_agents_enrolled(timeout) -> bool:
     """
     Wait for agents to be enrolled

deploy/test-environments/fleet_api/src/install_cnvm_integration.py

@@ -27,6 +27,7 @@
 CNVM_PACKAGE_POLICY = "../../../cloud/data/package_policy_cnvm_aws.json"
 CNVM_EXPECTED_AGENTS = 1
 CNVM_CLOUDFORMATION_CONFIG = "../../../cloudformation/config.json"
+CNVM_AGENT_TAGS = ["cft_version:CFT_VERSION", "cft_arn:arn:aws:cloudformation:.*"]
 
 cnvm_agent_policy_data = Path(__file__).parent / CNVM_AGENT_POLICY
 cnvm_pkg_policy_data = Path(__file__).parent / CNVM_PACKAGE_POLICY
@@ -60,7 +61,7 @@ def load_data() -> Tuple[Dict, Dict]:
         agent_policy_id=agent_policy_id,
     )
 
-    state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, CNVM_EXPECTED_AGENTS))
+    state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, CNVM_EXPECTED_AGENTS, CNVM_AGENT_TAGS))
 
     cloudformation_params = Munch()
     cloudformation_params.ENROLLMENT_TOKEN = get_enrollment_token(

deploy/test-environments/fleet_api/src/install_cspm_integration.py

@@ -76,7 +76,7 @@ def load_data() -> Tuple[Dict, Dict]:
         cspm_data=cspm_data,
     )
 
-    state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, CSPM_EXPECTED_AGENTS))
+    state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, CSPM_EXPECTED_AGENTS, []))
 
     manifest_params = Munch()
     manifest_params.enrollment_token = get_enrollment_token(

deploy/test-environments/fleet_api/src/install_kspm_eks_integration.py

@@ -71,7 +71,7 @@ def load_data() -> Tuple[Dict, Dict]:
         eks_data=eks_data,
     )
 
-    state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, KSPM_EKS_EXPECTED_AGENTS))
+    state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, KSPM_EKS_EXPECTED_AGENTS, []))
 
     manifest_params = Munch()
     manifest_params.enrollment_token = get_enrollment_token(

deploy/test-environments/fleet_api/src/install_kspm_unmanaged_integration.py

@@ -66,7 +66,7 @@ def load_data() -> Tuple[Dict, Dict]:
         agent_policy_id=agent_policy_id,
     )
 
-    state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, KSPM_UNMANAGED_EXPECTED_AGENTS))
+    state_manager.add_policy(PolicyState(agent_policy_id, package_policy_id, KSPM_UNMANAGED_EXPECTED_AGENTS, []))
 
     manifest_params = Munch()
     manifest_params.enrollment_token = get_enrollment_token(

deploy/test-environments/fleet_api/src/state_file_manager.py

@@ -27,7 +27,7 @@ class PolicyState:
     Class to represent a policy state.
     """
 
-    def __init__(self, agnt_policy_id: str, pkg_policy_id: str, expected_agents: int):
+    def __init__(self, agnt_policy_id: str, pkg_policy_id: str, expected_agents: int, expected_tags: list[str]):
         """
         Args:
             agnt_policy_id (str): ID of the agent policy.
@@ -37,6 +37,7 @@ def __init__(self, agnt_policy_id: str, pkg_policy_id: str, expected_agents: int
         self.agnt_policy_id = agnt_policy_id
         self.pkg_policy_id = pkg_policy_id
         self.expected_agents = expected_agents
+        self.expected_tags = expected_tags
 
 
 class StateFileManager: