opentelemetry-instrumentation: make it easier to use bootstrap with custom values (#2918)

Author: xrmx

opentelemetry-instrumentation/src/opentelemetry/instrumentation/bootstrap.py

@@ -22,12 +22,15 @@
     SubprocessError,
     check_call,
 )
+from typing import Optional
 
 from packaging.requirements import Requirement
 
 from opentelemetry.instrumentation.bootstrap_gen import (
-    default_instrumentations,
-    libraries,
+    default_instrumentations as gen_default_instrumentations,
+)
+from opentelemetry.instrumentation.bootstrap_gen import (
+    libraries as gen_libraries,
 )
 from opentelemetry.instrumentation.version import __version__
 from opentelemetry.util._importlib_metadata import (
@@ -75,7 +78,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 +116,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 +125,25 @@ 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()
+    _pip_check(libraries)
 
 
-def run() -> None:
+def run(
+    default_instrumentations: Optional[list] = None,
+    libraries: Optional[list] = None,
+) -> None:
     action_install = "install"
     action_requirements = "requirements"
 
@@ -163,8 +173,14 @@ def run() -> None:
     )
     args = parser.parse_args()
 
+    if libraries is None:
+        libraries = gen_libraries
+
+    if default_instrumentations is None:
+        default_instrumentations = gen_default_instrumentations
+
     cmd = {
         action_install: _run_install,
         action_requirements: _run_requirements,
     }[args.action]
-    cmd()
+    cmd(default_instrumentations, libraries)

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,15 @@ 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()
 
-    @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()
 
     @patch("sys.argv", ["bootstrap", "-a", "pipenv"])
     def test_run_unknown_cmd(self):
@@ -73,18 +76,44 @@ def test_run_unknown_cmd(self):
 
     @patch("sys.argv", ["bootstrap", "-a", "requirements"])
     def test_run_cmd_print(self):
+        self.pkg_patcher.start()
         with patch("sys.stdout", new=StringIO()) as fake_out:
             bootstrap.run()
             self.assertEqual(
                 fake_out.getvalue(),
                 "\n".join(self.installed_libraries) + "\n",
             )
+        self.pkg_patcher.stop()
 
     @patch("sys.argv", ["bootstrap", "-a", "install"])
     def test_run_cmd_install(self):
+        self.pkg_patcher.start()
         bootstrap.run()
         self.mock_pip_install.assert_has_calls(
             [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()
+        self.pkg_patcher.stop()
+
+    @patch("sys.argv", ["bootstrap", "-a", "install"])
+    def test_can_override_available_libraries(self):
+        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):
+        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()