From 2f5b40e00e4adb57c20896d484b23901e3c53278 Mon Sep 17 00:00:00 2001
From: Riccardo Magliocchetti <riccardo.magliocchetti@gmail.com>
Date: Fri, 20 Sep 2024 16:55:01 +0200
Subject: [PATCH] opentelemetry-instrumentation: make it easier to use
 bootstrap with custom values

Pass the available instrumentation libraries and the default libraries
from bootstrap_gen to run() (and down to all the functions) instead of
relying on the globals. This way it's easier to reuse run with custom
values.

While at it fix mock handling in tests.
---
 .../instrumentation/bootstrap.py              | 22 ++++++---
 .../tests/test_bootstrap.py                   | 49 +++++++++++++++----
 2 files changed, 53 insertions(+), 18 deletions(-)

diff --git a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/bootstrap.py b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/bootstrap.py
index 02926ea5c4..8b4bec1c27 100644
--- a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/bootstrap.py
+++ b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/bootstrap.py
@@ -75,7 +75,7 @@ def _sys_pip_install(package):
         print(error)
 
 
-def _pip_check():
+def _pip_check(libraries):
     """Ensures none of the instrumentations have dependency conflicts.
     Clean check reported as:
     'No broken requirements found.'
@@ -113,7 +113,7 @@ def _is_installed(req):
     return True
 
 
-def _find_installed_libraries():
+def _find_installed_libraries(default_instrumentations, libraries):
     for lib in default_instrumentations:
         yield lib
 
@@ -122,18 +122,24 @@ def _find_installed_libraries():
             yield lib["instrumentation"]
 
 
-def _run_requirements():
+def _run_requirements(default_instrumentations, libraries):
     logger.setLevel(logging.ERROR)
-    print("\n".join(_find_installed_libraries()))
+    print(
+        "\n".join(
+            _find_installed_libraries(default_instrumentations, libraries)
+        )
+    )
 
 
-def _run_install():
-    for lib in _find_installed_libraries():
+def _run_install(default_instrumentations, libraries):
+    for lib in _find_installed_libraries(default_instrumentations, libraries):
         _sys_pip_install(lib)
     _pip_check()
 
 
-def run() -> None:
+def run(
+    default_instrumentations=default_instrumentations, libraries=libraries
+) -> None:
     action_install = "install"
     action_requirements = "requirements"
 
@@ -167,4 +173,4 @@ def run() -> None:
         action_install: _run_install,
         action_requirements: _run_requirements,
     }[args.action]
-    cmd()
+    cmd(default_instrumentations, libraries)
diff --git a/opentelemetry-instrumentation/tests/test_bootstrap.py b/opentelemetry-instrumentation/tests/test_bootstrap.py
index 4807f0beb7..ab00effe7c 100644
--- a/opentelemetry-instrumentation/tests/test_bootstrap.py
+++ b/opentelemetry-instrumentation/tests/test_bootstrap.py
@@ -19,7 +19,10 @@
 from unittest.mock import call, patch
 
 from opentelemetry.instrumentation import bootstrap
-from opentelemetry.instrumentation.bootstrap_gen import libraries
+from opentelemetry.instrumentation.bootstrap_gen import (
+    default_instrumentations,
+    libraries,
+)
 
 
 def sample_packages(packages, rate):
@@ -56,15 +59,17 @@ def setUpClass(cls):
             "opentelemetry.instrumentation.bootstrap._pip_check",
         )
 
-        cls.pkg_patcher.start()
-        cls.mock_pip_install = cls.pip_install_patcher.start()
-        cls.mock_pip_check = cls.pip_check_patcher.start()
+    def setUp(self):
+        super().setUp()
+        self.mock_pip_check = self.pip_check_patcher.start()
+        self.mock_pip_install = self.pip_install_patcher.start()
+        self.pkg_patcher.start()
 
-    @classmethod
-    def tearDownClass(cls):
-        cls.pip_check_patcher.start()
-        cls.pip_install_patcher.start()
-        cls.pkg_patcher.stop()
+    def tearDown(self):
+        super().tearDown()
+        self.pip_check_patcher.stop()
+        self.pip_install_patcher.stop()
+        self.pkg_patcher.stop()
 
     @patch("sys.argv", ["bootstrap", "-a", "pipenv"])
     def test_run_unknown_cmd(self):
@@ -87,4 +92,28 @@ def test_run_cmd_install(self):
             [call(i) for i in self.installed_libraries],
             any_order=True,
         )
-        self.assertEqual(self.mock_pip_check.call_count, 1)
+        self.mock_pip_check.assert_called_once()
+
+    @patch("sys.argv", ["bootstrap", "-a", "install"])
+    def test_can_override_available_libraries(self):
+        self.pkg_patcher.stop()
+        bootstrap.run(libraries=[])
+        self.mock_pip_install.assert_has_calls(
+            [call(i) for i in default_instrumentations],
+            any_order=True,
+        )
+        self.mock_pip_check.assert_called_once()
+
+    @patch("sys.argv", ["bootstrap", "-a", "install"])
+    def test_can_override_available_default_instrumentations(self):
+        self.pkg_patcher.stop()
+        with patch(
+            "opentelemetry.instrumentation.bootstrap._is_installed",
+            return_value=True,
+        ):
+            bootstrap.run(default_instrumentations=[])
+        self.mock_pip_install.assert_has_calls(
+            [call(i) for i in self.installed_libraries],
+            any_order=True,
+        )
+        self.mock_pip_check.assert_called_once()