Add force flush before exit (#24)

* Add force flush

* Remove xdist method from test cov

* Updates following PR

* Restore accidentally deleted doc line!

Author: cpnat

src/pytest_opentelemetry/instrumentation.py

@@ -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)
 
@@ -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]:
@@ -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()

tests/test_sessions.py

@@ -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
@@ -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