From cc5f5b313779082a5383ba43e0e88f354174734f Mon Sep 17 00:00:00 2001
From: Romain Marcadier <romain@datadoghq.com>
Date: Tue, 2 Jul 2024 16:22:48 +0200
Subject: [PATCH] Use ddapm-test-agent from environment if available

---
 .github/workflows/tests.yml                   |  4 ++
 _integration-tests/utils/agent/agent.go       | 41 +++++++++++--------
 .../utils/agent/requirements.txt              |  1 +
 3 files changed, 30 insertions(+), 16 deletions(-)
 create mode 100644 _integration-tests/utils/agent/requirements.txt

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 334a835a..25f87b31 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -101,6 +101,10 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: '>=3.9 <3.13'
+          cache: pip
+          cache-dependency-path: _integration-tests/utils/agent/requirements.txt
+      - name: Install python dependencies
+        run: pip install -r _integration-tests/utils/agent/requirements.txt
       - name: Build orchestrion binary
         run: go build -cover -covermode=atomic -coverpkg="./..." -o="bin/orchestrion.exe" .
       - name: Run Integration Tests
diff --git a/_integration-tests/utils/agent/agent.go b/_integration-tests/utils/agent/agent.go
index a61f34fa..80470caf 100644
--- a/_integration-tests/utils/agent/agent.go
+++ b/_integration-tests/utils/agent/agent.go
@@ -42,32 +42,41 @@ func New(t *testing.T) (*MockAgent, error) {
 		agent MockAgent
 		err   error
 	)
-	if agent.virtualEnv, err = os.MkdirTemp("", "orchestrion-integ-venv-*"); err != nil {
-		return nil, err
-	}
-	t.Logf("Creating Python venv at %q...\n", agent.virtualEnv)
-	if err = exec.Command("python3", "-m", "venv", agent.virtualEnv).Run(); err != nil {
-		return nil, err
-	}
-	venvBin := filepath.Join(agent.virtualEnv, "bin")
-	if runtime.GOOS == "windows" {
-		venvBin = filepath.Join(agent.virtualEnv, "Scripts")
-	}
 
-	t.Logf("Installing ddapm-test-agent in venv...\n")
-	if err = exec.Command(filepath.Join(venvBin, "pip"), "install", "ddapm-test-agent").Run(); err != nil {
-		return nil, err
+	ddapmTestAgent, _ := exec.LookPath("ddapm-test-agent")
+	if ddapmTestAgent == "" {
+		t.Log("No ddapm-test-agent found in $PATH, installing into a python venv...")
+		if agent.virtualEnv, err = os.MkdirTemp("", "orchestrion-integ-venv-*"); err != nil {
+			return nil, err
+		}
+		t.Logf("Creating Python venv at %q...\n", agent.virtualEnv)
+		if err = exec.Command("python3", "-m", "venv", agent.virtualEnv).Run(); err != nil {
+			return nil, err
+		}
+		venvBin := filepath.Join(agent.virtualEnv, "bin")
+		if runtime.GOOS == "windows" {
+			venvBin = filepath.Join(agent.virtualEnv, "Scripts")
+		}
+
+		t.Logf("Installing requirements in venv...\n")
+		_, thisFile, _, _ := runtime.Caller(0)
+		thisDir := filepath.Dir(thisFile)
+		if err = exec.Command(filepath.Join(venvBin, "pip"), "install", "-r", filepath.Join(thisDir, "requirements.txt")).Run(); err != nil {
+			return nil, err
+		}
+
+		ddapmTestAgent = filepath.Join(venvBin, "ddapm-test-agent")
 	}
 
 	if agent.port, err = getFreePort(); err != nil {
 		return nil, err
 	}
-	t.Logf("Starting ddapm-test-agent on port %d\n", agent.port)
+	t.Logf("Starting %s on port %d\n", ddapmTestAgent, agent.port)
 	var ctx context.Context
 	ctx, agent.processCancel = context.WithCancel(context.Background())
 	agent.process = exec.CommandContext(
 		ctx,
-		filepath.Join(venvBin, "ddapm-test-agent"),
+		ddapmTestAgent,
 		fmt.Sprintf("--port=%d", agent.port),
 	)
 	if err = agent.process.Start(); err != nil {
diff --git a/_integration-tests/utils/agent/requirements.txt b/_integration-tests/utils/agent/requirements.txt
new file mode 100644
index 00000000..c6c293e0
--- /dev/null
+++ b/_integration-tests/utils/agent/requirements.txt
@@ -0,0 +1 @@
+ddapm-test-agent~=1.17.0