Skip to content

Commit

Permalink
test: add scenario for OTEL context forwarding in parallel bulk
Browse files Browse the repository at this point in the history
  • Loading branch information
claudinoac committed Jul 23, 2024
1 parent 87133b1 commit fbfd3d4
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 2 deletions.
2 changes: 1 addition & 1 deletion noxfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ def test(session):
session.run(*pytest_argv())


@nox.session(python=["3.7", "3.12"])
@nox.session(python=["3.7", "3.8", "3.12"])
def test_otel(session):
session.install(
".[dev]",
Expand Down
23 changes: 22 additions & 1 deletion test_elasticsearch/test_otel.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@
# under the License.

import os

from unittest import mock
from elasticsearch import Elasticsearch
from elasticsearch import helpers
import pytest

try:
Expand Down Expand Up @@ -95,3 +97,22 @@ def test_detailed_span():
"db.elasticsearch.cluster.name": "e9106fc68e3044f0b1475b04bf4ffd5f",
"db.elasticsearch.node.name": "instance-0000000001",
}


@mock.patch("elasticsearch._otel.OpenTelemetry.recover_parent_context")
@mock.patch("elasticsearch.helpers.actions._process_bulk_chunk_success")
@mock.patch("elasticsearch.Elasticsearch.bulk")
def test_forward_otel_context_to_subthreads(
_call_bulk_mock, _process_bulk_success_mock, _mock_otel_recv_context
):
tracer, memory_exporter = setup_tracing()
es_client = Elasticsearch("http://localhost:9200")
es_client._otel = OpenTelemetry(enabled=True, tracer=tracer)

_call_bulk_mock.return_value = mock.Mock()
actions = ({"x": i} for i in range(100))
list(
helpers.parallel_bulk(es_client, actions, chunk_size=4)
)
# Ensures that the OTEL context has been forwarded to all chunks
assert es_client._otel.recover_parent_context.call_count == 25

0 comments on commit fbfd3d4

Please sign in to comment.