opentelemetry-python-contrib/instrumentation/opentelemetry-instrumentation-pymongo/tests/test_pymongo.py

# 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 unittest import mock

from opentelemetry import trace as trace_api
from opentelemetry.instrumentation.pymongo import (
    CommandTracer,
    PymongoInstrumentor,
)
from opentelemetry.test.test_base import TestBase


class TestPymongo(TestBase):
    def setUp(self):
        super().setUp()
        self.tracer = self.tracer_provider.get_tracer(__name__)

    def test_pymongo_instrumentor(self):
        mock_register = mock.Mock()
        patch = mock.patch(
            "pymongo.monitoring.register", side_effect=mock_register
        )
        with patch:
            PymongoInstrumentor().instrument()

        self.assertTrue(mock_register.called)

    def test_started(self):
        command_attrs = {
            "filter": "filter",
            "sort": "sort",
            "limit": "limit",
            "pipeline": "pipeline",
            "command_name": "find",
        }
        command_tracer = CommandTracer(self.tracer)
        mock_event = MockEvent(
            command_attrs, ("test.com", "1234"), "test_request_id"
        )
        command_tracer.started(event=mock_event)
        # the memory exporter can't be used here because the span isn't ended
        # yet
        # pylint: disable=protected-access
        span = command_tracer._pop_span(mock_event)
        self.assertIs(span.kind, trace_api.SpanKind.CLIENT)
        self.assertEqual(span.name, "mongodb.command_name.find")
        self.assertEqual(span.attributes["component"], "mongodb")
        self.assertEqual(span.attributes["db.type"], "mongodb")
        self.assertEqual(span.attributes["db.instance"], "database_name")
        self.assertEqual(span.attributes["db.statement"], "command_name find")
        self.assertEqual(span.attributes["net.peer.name"], "test.com")
        self.assertEqual(span.attributes["net.peer.port"], "1234")
        self.assertEqual(
            span.attributes["db.mongo.operation_id"], "operation_id"
        )
        self.assertEqual(
            span.attributes["db.mongo.request_id"], "test_request_id"
        )

        self.assertEqual(span.attributes["db.mongo.filter"], "filter")
        self.assertEqual(span.attributes["db.mongo.sort"], "sort")
        self.assertEqual(span.attributes["db.mongo.limit"], "limit")
        self.assertEqual(span.attributes["db.mongo.pipeline"], "pipeline")

    def test_succeeded(self):
        mock_event = MockEvent({})
        command_tracer = CommandTracer(self.tracer)
        command_tracer.started(event=mock_event)
        command_tracer.succeeded(event=mock_event)
        spans_list = self.memory_exporter.get_finished_spans()
        self.assertEqual(len(spans_list), 1)
        span = spans_list[0]
        self.assertEqual(
            span.attributes["db.mongo.duration_micros"], "duration_micros"
        )
        self.assertIs(
            span.status.canonical_code, trace_api.status.StatusCanonicalCode.OK
        )
        self.assertEqual(span.status.description, "reply")
        self.assertIsNotNone(span.end_time)

    def test_failed(self):
        mock_event = MockEvent({})
        command_tracer = CommandTracer(self.tracer)
        command_tracer.started(event=mock_event)
        command_tracer.failed(event=mock_event)

        spans_list = self.memory_exporter.get_finished_spans()
        self.assertEqual(len(spans_list), 1)
        span = spans_list[0]

        self.assertEqual(
            span.attributes["db.mongo.duration_micros"], "duration_micros"
        )
        self.assertIs(
            span.status.canonical_code,
            trace_api.status.StatusCanonicalCode.UNKNOWN,
        )
        self.assertEqual(span.status.description, "failure")
        self.assertIsNotNone(span.end_time)

    def test_multiple_commands(self):
        first_mock_event = MockEvent({}, ("firstUrl", "123"), "first")
        second_mock_event = MockEvent({}, ("secondUrl", "456"), "second")
        command_tracer = CommandTracer(self.tracer)
        command_tracer.started(event=first_mock_event)
        command_tracer.started(event=second_mock_event)
        command_tracer.succeeded(event=first_mock_event)
        command_tracer.failed(event=second_mock_event)

        spans_list = self.memory_exporter.get_finished_spans()
        self.assertEqual(len(spans_list), 2)
        first_span = spans_list[0]
        second_span = spans_list[1]

        self.assertEqual(first_span.attributes["db.mongo.request_id"], "first")
        self.assertIs(
            first_span.status.canonical_code,
            trace_api.status.StatusCanonicalCode.OK,
        )
        self.assertEqual(
            second_span.attributes["db.mongo.request_id"], "second"
        )
        self.assertIs(
            second_span.status.canonical_code,
            trace_api.status.StatusCanonicalCode.UNKNOWN,
        )


class MockCommand:
    def __init__(self, command_attrs):
        self.command_attrs = command_attrs

    def get(self, key, default=""):
        return self.command_attrs.get(key, default)


class MockEvent:
    def __init__(self, command_attrs, connection_id=None, request_id=""):
        self.command = MockCommand(command_attrs)
        self.connection_id = connection_id
        self.request_id = request_id

    def __getattr__(self, item):
        return item