From b70d381efaed353b6d27e69cafaf48911e4e4a1d Mon Sep 17 00:00:00 2001
From: Joakim Recht <joakim@uber.com>
Date: Thu, 1 Sep 2016 16:39:22 +0200
Subject: [PATCH] Make it possible to specify a default rerun filter

---
 flaky/_flaky_plugin.py       | 18 +++++++++++++++---
 flaky/flaky_pytest_plugin.py |  8 ++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/flaky/_flaky_plugin.py b/flaky/_flaky_plugin.py
index 1adac9b..615e7da 100644
--- a/flaky/_flaky_plugin.py
+++ b/flaky/_flaky_plugin.py
@@ -367,6 +367,12 @@ def add_force_flaky_options(add_option):
                  "least this many times (unless the test has its own flaky "
                  "decorator)."
         )
+        add_option(
+            '--rerun-filter',
+            action='store',
+            dest='rerun_filter',
+            help='If --force-flaky is specified, use this as the rerun filter '
+                 'function to determine if a test should be retried.')
 
     def _add_flaky_report(self, stream):
         """
@@ -608,7 +614,7 @@ def _get_test_callable_name(test):
         raise NotImplementedError  # pragma: no cover
 
     @classmethod
-    def _make_test_flaky(cls, test, max_runs, min_passes):
+    def _make_test_flaky(cls, test, max_runs, min_passes, rerun_filter=None):
         """
         Make a given test flaky.
 
@@ -624,7 +630,13 @@ def _make_test_flaky(cls, test, max_runs, min_passes):
             The value of the FlakyNames.MIN_PASSES attribute to use.
         :type min_passes:
             `int`
-        """
-        attrib_dict = defaults.default_flaky_attributes(max_runs, min_passes)
+        :param rerun_filter:
+            Rerun filter function to use to determine if a test should
+            run again.
+        :type rerun_filter:
+            `function`
+        """
+        attrib_dict = defaults.default_flaky_attributes(
+            max_runs, min_passes, rerun_filter=rerun_filter)
         for attr, value in attrib_dict.items():
             cls._set_flaky_attribute(test, attr, value)
diff --git a/flaky/flaky_pytest_plugin.py b/flaky/flaky_pytest_plugin.py
index 11a13dc..032ebf8 100644
--- a/flaky/flaky_pytest_plugin.py
+++ b/flaky/flaky_pytest_plugin.py
@@ -6,6 +6,8 @@
 
 from flaky._flaky_plugin import _FlakyPlugin
 
+import sys
+
 
 class FlakyXdist(object):
 
@@ -34,6 +36,7 @@ class FlakyPlugin(_FlakyPlugin):
     max_runs = None
     min_passes = None
     config = None
+    rerun_filter = None
     _call_infos = {}
     _PYTEST_WHEN_CALL = 'call'
     _PYTEST_OUTCOME_PASSED = 'passed'
@@ -71,6 +74,7 @@ def pytest_runtest_protocol(self, item, nextitem):
                 item,
                 self.max_runs,
                 self.min_passes,
+                rerun_filter=self.rerun_filter,
             )
         original_call_and_report = self.runner.call_and_report
         self._call_infos[item] = {}
@@ -198,6 +202,10 @@ def pytest_configure(self, config):
         self.max_runs = config.option.max_runs
         self.min_passes = config.option.min_passes
         self.runner = config.pluginmanager.getplugin("runner")
+        if config.option.rerun_filter:
+            module, attr = config.option.rerun_filter.rsplit('.', 1)
+            __import__(module)
+            self.rerun_filter = getattr(sys.modules[module], attr)
         if config.pluginmanager.hasplugin('xdist'):
             config.pluginmanager.register(FlakyXdist(self), name='flaky.xdist')
             self.config = config