From 78dfa51a8b427a9e3b7caa857905a867c569b9bc Mon Sep 17 00:00:00 2001
From: mukund-ananthu <mukundananthu@google.com>
Date: Sat, 2 Nov 2024 19:38:43 +0000
Subject: [PATCH] fix: Add flaky decorator to test

* TO mitigate failures occurring due to non-deterministic order and
  number of OpenTelemetry spans being exported: https://github.com/open-telemetry/opentelemetry-python/issues/4210
---
 noxfile.py                                              | 4 +++-
 owlbot.py                                               | 1 +
 tests/unit/pubsub_v1/publisher/test_publisher_client.py | 6 ++++++
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/noxfile.py b/noxfile.py
index 1ba75e6b2..c810940e7 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -47,7 +47,9 @@
 UNIT_TEST_EXTERNAL_DEPENDENCIES: List[str] = []
 UNIT_TEST_LOCAL_DEPENDENCIES: List[str] = []
 UNIT_TEST_DEPENDENCIES: List[str] = []
-UNIT_TEST_EXTRAS: List[str] = []
+UNIT_TEST_EXTRAS: List[str] = [
+    "flaky",
+]
 UNIT_TEST_EXTRAS_BY_PYTHON: Dict[str, List[str]] = {}
 
 SYSTEM_TEST_PYTHON_VERSIONS: List[str] = ["3.12"]
diff --git a/owlbot.py b/owlbot.py
index 204b30ba5..94547b990 100644
--- a/owlbot.py
+++ b/owlbot.py
@@ -335,6 +335,7 @@
     microgenerator=True,
     samples=True,
     cov_level=100,
+    unit_test_extras=["flaky"],
     versions=gcp.common.detect_versions(path="./google", default_first=True),
     unit_test_python_versions=["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"],
     system_test_python_versions=["3.12"],
diff --git a/tests/unit/pubsub_v1/publisher/test_publisher_client.py b/tests/unit/pubsub_v1/publisher/test_publisher_client.py
index abc33f8cb..5facd34c6 100644
--- a/tests/unit/pubsub_v1/publisher/test_publisher_client.py
+++ b/tests/unit/pubsub_v1/publisher/test_publisher_client.py
@@ -28,6 +28,8 @@
 
 import pytest
 import time
+from flaky import flaky
+from typing import cast, Callable, Any, TypeVar
 
 from opentelemetry import trace
 from google.api_core import gapic_v1
@@ -48,6 +50,9 @@
     PublishMessageWrapper,
 )
 
+C = TypeVar("C", bound=Callable[..., Any])
+typed_flaky = cast(Callable[[C], C], flaky(max_runs=5, min_passes=1))
+
 
 def _assert_retries_equal(retry, retry2):
     # Retry instances cannot be directly compared, because their predicates are
@@ -264,6 +269,7 @@ def test_opentelemetry_publisher_batching_exception(
     sys.version_info < (3, 8),
     reason="Open Telemetry not supported below Python version 3.8",
 )
+@typed_flaky
 def test_opentelemetry_flow_control_exception(creds, span_exporter):
     publisher_options = types.PublisherOptions(
         flow_control=types.PublishFlowControl(