From 8940b5949aa0acce5cd17bd1450801d3c84fc07e Mon Sep 17 00:00:00 2001
From: Michael Adkins <michael@prefect.io>
Date: Tue, 21 Dec 2021 12:44:36 -0600
Subject: [PATCH] Add service markers to avoid slow tests with pytest

---
 tests/conftest.py | 45 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/tests/conftest.py b/tests/conftest.py
index 3b55d878d285..25b5391a628c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -10,6 +10,35 @@
 from .fixtures.database import *
 
 
+def pytest_configure(config):
+    config.addinivalue_line(
+        "markers", "service(arg) a service integration test. For example, 'docker'."
+    )
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--service",
+        action="append",
+        metavar="SERVICE",
+        default=[],
+        help="include service integration tests for SERVICE.",
+    )
+    parser.addoption(
+        "--all-services",
+        action="store_true",
+        default=False,
+        help="include all service integration tests",
+    )
+
+
+def pytest_runtest_setup(item):
+    envnames = [mark.args[0] for mark in item.iter_markers(name="env")]
+    if envnames:
+        if item.config.getoption("-E") not in envnames:
+            pytest.skip("test requires env in {!r}".format(envnames))
+
+
 def pytest_collection_modifyitems(session, config, items):
     """
     Modify all tests to automatically and transparently support asyncio
@@ -21,6 +50,22 @@ def pytest_collection_modifyitems(session, config, items):
         ):
             item.add_marker(pytest.mark.asyncio)
 
+    if config.getoption("--all-services"):
+        # Do not skip any service tests
+        return
+
+    run_services = set(config.getoption("--service"))
+    for item in items:
+        item_services = {mark.args[0] for mark in item.iter_markers(name="service")}
+        missing_services = item_services.difference(run_services)
+        if missing_services:
+            item.add_marker(
+                pytest.mark.skip(
+                    f"Requires services: {', '.join(missing_services)}. "
+                    "Use '--service NAME' to include test."
+                )
+            )
+
 
 @pytest.fixture(scope="session")
 def event_loop(request):