Add retry decorator & use it in flakey test

Author: hmstepanek

tests/agent_streaming/test_infinite_tracing.py

@@ -17,7 +17,7 @@
 
 import pytest
 from testing_support.fixtures import override_generic_settings
-from testing_support.util import conditional_decorator
+from testing_support.util import conditional_decorator, retry
 from testing_support.validators.validate_metric_payload import validate_metric_payload
 
 from newrelic.common.streaming_utils import StreamBuffer
@@ -333,6 +333,7 @@ def connect_complete():
 @conditional_decorator(
     condition=six.PY2, decorator=pytest.mark.xfail(reason="Test frequently times out on Py2.", strict=False)
 )
+@retry(attempts=5, wait=2)  # This test is flakey so add a retry.
 def test_no_data_loss_on_reconnect(mock_grpc_server, app, buffer_empty_event, batching, spans_processed_event):
     """
     Test for data loss when channel is closed by the server while waiting for more data in a request iterator.
@@ -350,10 +351,6 @@ def test_no_data_loss_on_reconnect(mock_grpc_server, app, buffer_empty_event, ba
     Relevant GitHub issue: https://github.com/grpc/grpc/issues/29110
     """
 
-    # The nonbatching test frequently fails on CI so skip it for now.
-    if batching == "nonbatching":
-        pytest.skip("Flakey test")
-
     terminating_span = Span(
         intrinsics={"wait_then_ok": AttributeValue(string_value="OK")}, agent_attributes={}, user_attributes={}
     )

tests/testing_support/util.py

@@ -14,11 +14,15 @@
 
 import re
 import socket
+import time
+from functools import wraps
+
 
 def _to_int(version_str):
-    m = re.match(r'\d+', version_str)
+    m = re.match(r"\d+", version_str)
     return int(m.group(0)) if m else 0
 
+
 def version2tuple(version_str):
     """Convert version, even if it contains non-numeric chars.
 
@@ -27,32 +31,55 @@ def version2tuple(version_str):
 
     """
 
-    parts = version_str.split('.')[:2]
+    parts = version_str.split(".")[:2]
     return tuple(map(_to_int, parts))
 
+
 def instance_hostname(hostname):
-    if hostname == 'localhost' or hostname == "127.0.0.1":
+    if hostname == "localhost" or hostname == "127.0.0.1":
         hostname = socket.gethostname()
     return hostname
 
+
 def get_open_port():
     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     s.bind(("", 0))
     port = s.getsockname()[1]
     s.close()
     return port
 
+
 def conditional_decorator(condition, decorator):
     """Applies a decorator if the condition is true. Accepts 0 argument callables for the condition."""
+
     def _conditional_decorator(func):
         if callable(condition):
             condition_eval = condition()
         else:
             condition_eval = condition
-        
+
         if condition_eval:
             return decorator(func)
         else:
             return func
 
     return _conditional_decorator
+
+
+def retry(attempts=5, wait=5):
+    def decorator(test_func):
+        @wraps(test_func)
+        def wrapper(*args, **kwargs):
+            retry_count = 1
+            while retry_count < attempts:
+                try:
+                    return test_func(*args, **kwargs)
+                except AssertionError as assert_error:
+                    time.sleep(wait)
+                    retry_count += 1
+            # Preserve original traceback in case assertion fails.
+            return test_func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator