test(pytest): Refactor test grouping and allow selecting only snuba tests (#20915)

This refactors test grouping so that we deselect inside of `conftest.py` in the
`pytest_collection_modifyitems` hook. Previously, we would mark tests with the group, and at run-time, select the group to run. With this change, we can clean up a bunch of conditionals in our Makefile, and allows any task that uses pytest to use grouping without any further Makefile changes.

Also adds option to only select Snuba tests, as I need this for a future PR, although I could be convinced to move this into a new PR.

Makefile

@@ -126,12 +126,7 @@ fetch-release-registry:
 
 run-acceptance:
 	@echo "--> Running acceptance tests"
-ifndef TEST_GROUP
 	py.test tests/acceptance --cov . --cov-report="xml:.artifacts/acceptance.coverage.xml" --junit-xml=".artifacts/acceptance.junit.xml" --html=".artifacts/acceptance.pytest.html" --self-contained-html
-else
-	py.test tests/acceptance -m group_$(TEST_GROUP) --cov . --cov-report="xml:.artifacts/acceptance.coverage.xml" --junit-xml=".artifacts/acceptance.junit.xml" --html=".artifacts/acceptance.pytest.html" --self-contained-html
-endif
-
 	@echo ""
 
 test-cli:
@@ -164,20 +159,12 @@ test-js-ci: node-version-check
 test-python:
 	@echo "--> Running Python tests"
 	# This gets called by getsentry
-ifndef TEST_GROUP
 	py.test tests/integration tests/sentry
-else
-	py.test tests/integration tests/sentry -m group_$(TEST_GROUP)
-endif
 
 test-python-ci:
 	make build-platform-assets
 	@echo "--> Running CI Python tests"
-ifndef TEST_GROUP
 	py.test tests/integration tests/sentry --cov . --cov-report="xml:.artifacts/python.coverage.xml" --junit-xml=".artifacts/python.junit.xml" || exit 1
-else
-	py.test tests/integration tests/sentry -m group_$(TEST_GROUP) --cov . --cov-report="xml:.artifacts/python.coverage.xml" --junit-xml=".artifacts/python.junit.xml" || exit 1
-endif
 	@echo ""
 
 test-snuba:
@@ -199,12 +186,7 @@ test-plugins:
 	@echo "--> Building static assets"
 	@$(WEBPACK) --display errors-only
 	@echo "--> Running plugin tests"
-
-ifndef TEST_GROUP
 	py.test tests/sentry_plugins -vv --cov . --cov-report="xml:.artifacts/plugins.coverage.xml" --junit-xml=".artifacts/plugins.junit.xml" || exit 1
-else
-	py.test tests/sentry_plugins -m group_$(TEST_GROUP) -vv --cov . --cov-report="xml:.artifacts/plugins.coverage.xml" --junit-xml=".artifacts/plugins.junit.xml" || exit 1
-endif
 	@echo ""
 
 test-relay-integration:

conftest.py

@@ -2,10 +2,6 @@ from __future__ import absolute_import
 
 import os
 import sys
-from hashlib import md5
-
-import six
-import pytest
 
 pytest_plugins = ["sentry.utils.pytest"]
 
@@ -48,15 +44,3 @@ def install_sentry_plugins():
     settings.GITHUB_API_SECRET = "123"
     # this isn't the real secret
     settings.SENTRY_OPTIONS["github.integration-hook-secret"] = "b3002c3e321d4b7880360d397db2ccfd"
-
-
-def pytest_collection_modifyitems(config, items):
-    for item in items:
-        total_groups = int(os.environ.get("TOTAL_TEST_GROUPS", 1))
-        # TODO(joshuarli): six 1.12.0 adds ensure_binary: six.ensure_binary(item.location[0])
-        group_num = (
-            int(md5(six.text_type(item.location[0]).encode("utf-8")).hexdigest(), 16) % total_groups
-        )
-        marker = "group_%s" % group_num
-        config.addinivalue_line("markers", marker)
-        item.add_marker(getattr(pytest.mark, marker))

src/sentry/utils/pytest/sentry.py

@@ -2,10 +2,12 @@ from __future__ import absolute_import
 
 from sentry.utils.compat import mock
 import os
+from hashlib import md5
 
 from django.conf import settings
 from sentry_sdk import Hub
 
+import six
 
 TEST_ROOT = os.path.normpath(
     os.path.join(os.path.dirname(__file__), os.pardir, os.pardir, os.pardir, os.pardir, "tests")
@@ -263,3 +265,57 @@ def pytest_runtest_teardown(item):
         model.objects.clear_local_cache()
 
     Hub.main.bind_client(None)
+
+
+def pytest_collection_modifyitems(config, items):
+    """
+    After collection, we need to:
+
+    - Filter tests that subclass SnubaTestCase as tests in `tests/acceptance` are not being marked as `snuba`
+    - Select tests based on group and group strategy
+
+    """
+
+    total_groups = int(os.environ.get("TOTAL_TEST_GROUPS", 1))
+    current_group = int(os.environ.get("TEST_GROUP", 0))
+    grouping_strategy = os.environ.get("TEST_GROUP_STRATEGY", "file")
+
+    accepted, keep, discard = [], [], []
+
+    for index, item in enumerate(items):
+        # XXX: For some reason tests in `tests/acceptance` are not being
+        # marked as snuba, so deselect test cases not a subclass of SnubaTestCase
+        if os.environ.get("RUN_SNUBA_TESTS_ONLY"):
+            from sentry.testutils import SnubaTestCase
+            import inspect
+
+            if inspect.isclass(item.cls) and not issubclass(item.cls, SnubaTestCase):
+                # No need to group if we are deselecting this
+                discard.append(item)
+                continue
+            accepted.append(item)
+        else:
+            accepted.append(item)
+
+        # In the case where we group by round robin (e.g. TEST_GROUP_STRATEGY is not `file`),
+        # we want to only include items in `accepted` list
+
+        # TODO(joshuarli): six 1.12.0 adds ensure_binary: six.ensure_binary(item.location[0])
+        item_to_group = (
+            int(md5(six.text_type(item.location[0]).encode("utf-8")).hexdigest(), 16)
+            if grouping_strategy == "file"
+            else len(accepted) - 1
+        )
+
+        # Split tests in different groups
+        group_num = item_to_group % total_groups
+
+        if group_num == current_group:
+            keep.append(item)
+        else:
+            discard.append(item)
+
+    # This only needs to be done if there are items to be de-selected
+    if len(discard) > 0:
+        items[:] = keep
+        config.hook.pytest_deselected(items=discard)