From 1d0f982e95cb58a18d99d7b47447dd17bbc10ae0 Mon Sep 17 00:00:00 2001
From: Harshil <ilovebhagwan@gmail.com>
Date: Wed, 21 Apr 2021 14:29:24 +0530
Subject: [PATCH] Don't fail tests if they've exceeded the timeout

---
 requirements-dev.txt |  1 -
 setup.cfg            |  2 ++
 tests/conftest.py    | 41 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/requirements-dev.txt b/requirements-dev.txt
index 28a12e942..b3b33b902 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -13,5 +13,4 @@ pytest==6.2.2
 
 flaky
 beautifulsoup4
-pytest-timeout
 wheel
diff --git a/setup.cfg b/setup.cfg
index dd6d88012..51946e5f7 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -23,6 +23,8 @@ disable = C0330,R0801,R0913,R0904,R0903,R0902,W0511,C0116,C0115,W0703,R0914,R091
 
 [tool:pytest]
 testpaths = tests
+markers =
+    timeout: xfail a test if exceeds the timeout specified.
 addopts = --no-success-flaky-report -rsxX
 filterwarnings =
     error
diff --git a/tests/conftest.py b/tests/conftest.py
index 621bc9af2..733f600c2 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -21,6 +21,7 @@ import functools
 import inspect
 import os
 import re
+import threading
 from collections import defaultdict
 from queue import Queue
 from threading import Thread, Event
@@ -339,6 +340,46 @@ def timezone(tzinfo):
     return tzinfo
 
 
+timed_out = False
+
+
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_call(item):
+    marker = item.get_closest_marker("timeout")  # Check if function has timeout marker
+    max_time = 10  # Default timeout if timeout arg not specified
+    if marker:
+        max_time = marker.args[0] if marker.args else 10  # Get timeout arg
+        timeout_setup(item, max_time)  # Set and start timer
+    yield  # This executes the test function and later the code after this is run.
+    if marker:
+        timeout_teardown(item)  # Cancel timer if timeout didn't exceed.
+        if timed_out:  # xfail test if exceeded!
+            pytest.xfail(reason=f"Timeout of {max_time}s exceeded.")
+
+
+def timeout_setup(item, max_time):
+    global timed_out
+
+    def xfail_test():
+        global timed_out
+        timed_out = True
+
+    timer = threading.Timer(max_time, xfail_test)
+
+    def cancel_timer():
+        timer.cancel()
+        timer.join()
+
+    item.cancel_timeout = cancel_timer
+    timer.start()
+
+
+def timeout_teardown(item):
+    cancel = getattr(item, "cancel_timeout", None)
+    if cancel:
+        cancel()
+
+
 def expect_bad_request(func, message, reason):
     """
     Wrapper for testing bot functions expected to result in an :class:`telegram.error.BadRequest`.