Add tests around the psycopg2 and psycopg2-binary packages to make sure auto instrumentation works for both

Author: rjduffner

instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_instrumentation.py

@@ -0,0 +1,183 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# type: ignore
+
+
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from importlib.metadata import PackageNotFoundError
+from unittest import TestCase
+from unittest.mock import call, patch, Mock
+
+from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
+from opentelemetry.instrumentation.psycopg2.package import (
+    _instruments,
+    _instruments_psycopg2,
+    _instruments_psycopg2_binary,
+)
+from opentelemetry.instrumentation.auto_instrumentation._load import (
+    _load_instrumentors,
+)
+
+
+
+class TestPsycopg2InstrumentationDependencies(TestCase):
+    """
+    The psycopg2 instrumentation is special in that it can instrument
+    multiple packages. (psycopg2 and psycopg2-binary)
+    The tests below are to verify that the correct package is instrumented
+    depending on which package is installed.
+    Including the cases where both are installed and neither is installed.
+    These tests were added after we found a bug where the auto instrumentation
+    path attempted to use a different dependency check than the manual
+    instrumentation path. This caused the auto instrumentation to fail.
+
+    """
+
+    @patch("opentelemetry.instrumentation.psycopg2.distribution")
+    def test_instrumentation_dependencies_psycopg2_installed(
+        self, mock_distribution
+    ) -> None:
+        instrumentation = Psycopg2Instrumentor()
+
+        def _distribution(name):
+            if name == "psycopg2":
+                return None
+            raise PackageNotFoundError
+
+        mock_distribution.side_effect = _distribution
+        package_to_instrument = instrumentation.instrumentation_dependencies()
+
+        self.assertEqual(mock_distribution.call_count, 1)
+        self.assertEqual(
+            mock_distribution.mock_calls,
+            [
+                call("psycopg2"),
+            ],
+        )
+        self.assertEqual(package_to_instrument, (_instruments_psycopg2,))
+
+    @patch("opentelemetry.instrumentation.psycopg2.distribution")
+    def test_instrumentation_dependencies_psycopg2_binary_installed(
+        self, mock_distribution
+    ) -> None:
+        instrumentation = Psycopg2Instrumentor()
+
+        def _distribution(name):
+            if name == "psycopg2-binary":
+                return None
+            raise PackageNotFoundError
+
+        mock_distribution.side_effect = _distribution
+        package_to_instrument = instrumentation.instrumentation_dependencies()
+
+        self.assertEqual(mock_distribution.call_count, 2)
+        self.assertEqual(
+            mock_distribution.mock_calls,
+            [
+                call("psycopg2"),
+                call("psycopg2-binary"),
+            ]
+        )
+        self.assertEqual(
+            package_to_instrument, (_instruments_psycopg2_binary,)
+        )
+
+    @patch("opentelemetry.instrumentation.psycopg2.distribution")
+    def test_instrumentation_dependencies_both_installed(
+        self, mock_distribution
+    ) -> None:
+        instrumentation = Psycopg2Instrumentor()
+
+        def _distribution(name):
+            # Function raises PackageNotFoundError
+            # if name is not in the list. We will
+            # not raise it for both names
+            return None
+
+        mock_distribution.side_effect = _distribution
+        package_to_instrument = instrumentation.instrumentation_dependencies()
+
+        self.assertEqual(mock_distribution.call_count, 1)
+        self.assertEqual(
+            mock_distribution.mock_calls, [call("psycopg2")]
+        )
+        self.assertEqual(
+            package_to_instrument, (_instruments_psycopg2,)
+        )
+
+    @patch("opentelemetry.instrumentation.psycopg2.distribution")
+    def test_instrumentation_dependencies_none_installed(
+        self, mock_distribution
+    ) -> None:
+        instrumentation = Psycopg2Instrumentor()
+
+        def _distribution(name):
+            # Function raises PackageNotFoundError
+            # if name is not in the list. We will
+            # raise it for both names to simulate
+            # neither being installed
+            raise PackageNotFoundError
+
+        mock_distribution.side_effect = _distribution
+        package_to_instrument = instrumentation.instrumentation_dependencies()
+
+        assert mock_distribution.call_count == 2
+        assert mock_distribution.mock_calls == [
+            call("psycopg2"),
+            call("psycopg2-binary"),
+        ]
+        assert package_to_instrument == _instruments
+
+        self.assertEqual(mock_distribution.call_count, 2)
+        self.assertEqual(
+            mock_distribution.mock_calls,
+            [
+                call("psycopg2"),
+                call("psycopg2-binary"),
+            ],
+        )
+        self.assertEqual(package_to_instrument, _instruments)
+
+    # This test is to verify that the auto instrumentation path
+    # will auto instrument psycopg2 or psycopg2-binary is installed.
+    # Note there is only one test here but it is run twice in tox
+    # once with the psycopg2 package installed and once with
+    # psycopg2-binary installed.
+    @patch("opentelemetry.instrumentation.auto_instrumentation._load._logger")
+    def test_instruments_with_psycopg2_installed(self, mock_logger):
+
+        def _instrumentation_loaded_successfully_call():
+            return call("Instrumented %s", "psycopg2")
+
+        mock_distro = Mock()
+        mock_distro.load_instrumentor.return_value = None
+        _load_instrumentors(mock_distro)
+        self.assertEqual(len(mock_distro.load_instrumentor.call_args_list), 1)
+        (ep,) = mock_distro.load_instrumentor.call_args.args
+        self.assertEqual(ep.name, "psycopg2")
+        mock_logger.debug.assert_has_calls(
+            [_instrumentation_loaded_successfully_call()]
+        )