diff --git a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/bootstrap.py b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/bootstrap.py index 02926ea5c4..cc0ac68f1c 100644 --- a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/bootstrap.py +++ b/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) diff --git a/opentelemetry-instrumentation/tests/test_bootstrap.py b/opentelemetry-instrumentation/tests/test_bootstrap.py index 4807f0beb7..0ded8d37b1 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,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()