From fbf9c48ab75012ffe898c0bf395ffb57cac82341 Mon Sep 17 00:00:00 2001
From: William Yu <william.yu@datadoghq.com>
Date: Thu, 20 Jun 2024 16:23:37 -0400
Subject: [PATCH] add manual process discovery and container checks for k8s e2e

---
 test/new-e2e/tests/process/k8s_test.go | 30 +++++++++++++++++++++++++-
 test/new-e2e/tests/process/testing.go  | 24 ++++++++++++++++++---
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/test/new-e2e/tests/process/k8s_test.go b/test/new-e2e/tests/process/k8s_test.go
index 597574b465a63b..4c72573f361747 100644
--- a/test/new-e2e/tests/process/k8s_test.go
+++ b/test/new-e2e/tests/process/k8s_test.go
@@ -57,6 +57,7 @@ type K8sSuite struct {
 }
 
 func TestK8sTestSuite(t *testing.T) {
+	t.Parallel()
 	helmValues, err := createHelmValues(helmConfig{
 		ProcessAgentEnabled: true,
 		ProcessCollection:   true,
@@ -79,7 +80,7 @@ func (s *K8sSuite) TestManualProcessCheck() {
 	agent := getAgentPod(s.T(), s.Env().KubernetesCluster.Client())
 
 	// The log level needs to be overridden as the pod has an ENV var set.
-	// This is to so we get just json back from the check
+	// This is so we get just json back from the check
 	stdout, stderr, err := s.Env().KubernetesCluster.KubernetesClient.
 		PodExec(agent.Namespace, agent.Name, "process-agent",
 			[]string{"bash", "-c", "DD_LOG_LEVEL=OFF process-agent check process -w 5s --json"})
@@ -89,6 +90,33 @@ func (s *K8sSuite) TestManualProcessCheck() {
 	assertManualProcessCheck(s.T(), stdout, false, "stress-ng-cpu [run]", "stress-ng")
 }
 
+func (s *K8sSuite) TestManualProcessDiscoveryCheck() {
+	agent := getAgentPod(s.T(), s.Env().KubernetesCluster.Client())
+	// The log level needs to be overridden as the pod has an ENV var set.
+	// This is so we get just json back from the check
+	stdout, stderr, err := s.Env().KubernetesCluster.KubernetesClient.
+		PodExec(agent.Namespace, agent.Name, "process-agent",
+			[]string{"bash", "-c", "DD_LOG_LEVEL=OFF process-agent check process_discovery -w 5s --json"})
+	assert.NoError(s.T(), err)
+	assert.Empty(s.T(), stderr)
+
+	assertManualProcessDiscoveryCheck(s.T(), stdout, "stress-ng-cpu [run]")
+}
+
+func (s *K8sSuite) TestManualContainerCheck() {
+	agent := getAgentPod(s.T(), s.Env().KubernetesCluster.Client())
+
+	// The log level needs to be overridden as the pod has an ENV var set.
+	// This is so we get just json back from the check
+	stdout, stderr, err := s.Env().KubernetesCluster.KubernetesClient.
+		PodExec(agent.Namespace, agent.Name, "process-agent",
+			[]string{"bash", "-c", "DD_LOG_LEVEL=OFF process-agent check container -w 5s --json"})
+	assert.NoError(s.T(), err)
+	assert.Empty(s.T(), stderr)
+
+	assertManualContainerCheck(s.T(), stdout, "stress-ng")
+}
+
 func getAgentPod(t *testing.T, client kubeClient.Interface) corev1.Pod {
 	res, err := client.CoreV1().Pods("datadog").
 		List(context.Background(), v1.ListOptions{LabelSelector: "app=dda-linux-datadog"})
diff --git a/test/new-e2e/tests/process/testing.go b/test/new-e2e/tests/process/testing.go
index 110a1c98a6ae0f..0e49cf56aff8d9 100644
--- a/test/new-e2e/tests/process/testing.go
+++ b/test/new-e2e/tests/process/testing.go
@@ -250,8 +250,7 @@ func assertManualProcessCheck(t *testing.T, check string, withIOStats bool, proc
 	}()
 
 	var checkOutput struct {
-		Processes  []*agentmodel.Process   `json:"processes"`
-		Containers []*agentmodel.Container `json:"containers"`
+		Processes []*agentmodel.Process `json:"processes"`
 	}
 
 	err := json.Unmarshal([]byte(check), &checkOutput)
@@ -275,8 +274,27 @@ func assertManualProcessCheck(t *testing.T, check string, withIOStats bool, proc
 		assert.Truef(t, hasIOStats, "Missing IOStats: %+v", procs)
 	}
 
+	assertManualContainerCheck(t, check, expectedContainers...)
+}
+
+// assertManualContainerCheck asserts that the given container is collected from a manual container check
+func assertManualContainerCheck(t *testing.T, check string, expectedContainers ...string) {
+	defer func() {
+		if t.Failed() {
+			t.Logf("Check output:\n%s\n", check)
+		}
+	}()
+
+	var checkOutput struct {
+		Containers []*agentmodel.Container `json:"containers"`
+	}
+
+	err := json.Unmarshal([]byte(check), &checkOutput)
+	require.NoError(t, err, "failed to unmarshal process check output")
+
 	for _, container := range expectedContainers {
-		assert.Truef(t, findContainer(container, checkOutput.Containers), "%s container not found in %+v", container, checkOutput.Containers)
+		assert.Truef(t, findContainer(container, checkOutput.Containers),
+			"%s container not found in %+v", container, checkOutput.Containers)
 	}
 }