Skip to content

Commit

Permalink
Add force flush before exit (#24)
Browse files Browse the repository at this point in the history
* Add force flush

* Remove xdist method from test cov

* Updates following PR

* Restore accidentally deleted doc line!
  • Loading branch information
cpnat authored Mar 15, 2023
1 parent 08b58c1 commit 7f0246c
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 0 deletions.
16 changes: 16 additions & 0 deletions src/pytest_opentelemetry/instrumentation.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,17 @@ def get_trace_parent(cls, config: Config) -> Optional[Context]:

return None

@classmethod
def try_force_flush(cls) -> bool:
provider = trace.get_tracer_provider()

# Not all providers (e.g. ProxyTraceProvider) implement force flush
if hasattr(provider, 'force_flush'):
provider.force_flush()
return True
else:
return False

def pytest_configure(self, config: Config) -> None:
self.trace_parent = self.get_trace_parent(config)

Expand All @@ -75,7 +86,9 @@ def pytest_sessionfinish(self, session: Session) -> None:
self.session_span.set_status(
StatusCode.ERROR if self.has_error else StatusCode.OK
)

self.session_span.end()
self.try_force_flush()

@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_protocol(self, item: Item) -> Generator[None, None, None]:
Expand Down Expand Up @@ -159,3 +172,6 @@ def pytest_configure(self, config: Config) -> None:
def pytest_configure_node(self, node: WorkerController) -> None: # pragma: no cover
with trace.use_span(self.session_span, end_on_exit=False):
propagate.inject(node.workerinput)

def pytest_xdist_node_collection_finished(node, ids): # pragma: no cover
super().try_force_flush()
20 changes: 20 additions & 0 deletions tests/test_sessions.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import os
from contextlib import contextmanager
from typing import Dict, Generator, Optional
from unittest.mock import Mock, patch

from _pytest.pytester import Pytester
from opentelemetry import trace
Expand Down Expand Up @@ -184,3 +185,22 @@ def test_two():
'00-1234567890abcdef1234567890abcdef-fedcba0987654321-01',
)
result.assert_outcomes(passed=2)


@patch.object(trace, 'get_tracer_provider')
def test_force_flush_with_supported_provider(mock_get_tracer_provider):
provider = Mock()
provider.force_flush = Mock(return_value=None)
mock_get_tracer_provider.return_value = provider

for plugin in OpenTelemetryPlugin, XdistOpenTelemetryPlugin:
assert plugin.try_force_flush() is True


@patch.object(trace, 'get_tracer_provider')
def test_force_flush_with_unsupported_provider(mock_get_tracer_provider):
provider = Mock(spec=trace.ProxyTracerProvider)
mock_get_tracer_provider.return_value = provider

for plugin in OpenTelemetryPlugin, XdistOpenTelemetryPlugin:
assert plugin.try_force_flush() is False

0 comments on commit 7f0246c

Please sign in to comment.