Partial test runner

partial_runner.py

import unittest
from unittest import suite
from unittest.suite import BaseTestSuite

from django.test.runner import DiscoverRunner
from getenv import env


def test_key(suite):
    if isinstance(suite, BaseTestSuite):
        suite._tests = sorted(suite._tests, key=test_key)
    return repr(suite)


class PartialTestLoader(unittest.TestLoader):
    def __init__(self) -> None:
        super().__init__()
        self.num_groups = env('CI_NODE_TOTAL', 1)
        self.test_group = env('CI_NODE_INDEX', 1)

    def discover(self, *args, **kwargs) -> unittest.suite.TestSuite:
        tests: suite.TestSuite = super().discover(*args, **kwargs)
        if self.num_groups > 1:
            # Sort tests to ensure that each runner has a consistent set of tests with no overlap and no omitted tests
            test_list = sorted(tests._tests, key=test_key)
            # Slice the tests to only return tests from the selected group
            # Subtract 0 because slices are 0-based
            test_list = test_list[self.test_group - 1::self.num_groups]
            tests._tests = test_list
        return tests


class TestRunner(DiscoverRunner):
    test_loader = PartialTestLoader