Fix grouping edge cases

pytest_split_tests/__init__.py

@@ -1,48 +1,70 @@
 # -*- coding: utf-8 -*-
 import json
-import math
 from random import Random
 
 from _pytest.config import create_terminal_writer
 
 import pytest
 
 
-def get_group_size(total_items, total_groups):
-    """Return the group size."""
-    return int(math.ceil(float(total_items) / total_groups))
+def get_group(items, total_groups, group_id):
+    """Get the items from the passed in group based on group size."""
+    total_items = len(items)
+    base_group_size = total_items // total_groups
+    remainder = total_items % total_groups
+    start = 0
+    for i in range(1, group_id + 1):
+        if start >= len(items) or start < 0:
+            raise ValueError("Invalid test-group argument. Seems to be too high")
 
+        this_group_size = base_group_size + (1 if i <= remainder else 0)
 
-def get_group(items, group_size, group_id):
-    """Get the items from the passed in group based on group size."""
-    start = group_size * (group_id - 1)
-    end = start + group_size
+        if i == group_id:
+            end = start + this_group_size
+            return items[start:end]
 
-    if start >= len(items) or start < 0:
-        raise ValueError("Invalid test-group argument")
+        start += this_group_size
 
-    return items[start:end]
+    raise ValueError("Invalid test-group argument. Seems to be too low")
 
 
 def pytest_addoption(parser):
-    group = parser.getgroup('split your tests into evenly sized groups and run them')
-    group.addoption('--test-group-count', dest='test-group-count', type=int,
-                    help='The number of groups to split the tests into')
-    group.addoption('--test-group', dest='test-group', type=int,
-                    help='The group of tests that should be executed')
-    group.addoption('--test-group-random-seed', dest='random-seed', type=int, default=False,
-                    help='Integer to seed pseudo-random test selection')
-    group.addoption('--test-group-prescheduled', dest='prescheduled', type=str, default=None,
-                    help='Path to JSON file containing which tests to run.')
+    group = parser.getgroup("split your tests into evenly sized groups and run them")
+    group.addoption(
+        "--test-group-count",
+        dest="test-group-count",
+        type=int,
+        help="The number of groups to split the tests into",
+    )
+    group.addoption(
+        "--test-group",
+        dest="test-group",
+        type=int,
+        help="The group of tests that should be executed",
+    )
+    group.addoption(
+        "--test-group-random-seed",
+        dest="random-seed",
+        type=int,
+        default=False,
+        help="Integer to seed pseudo-random test selection",
+    )
+    group.addoption(
+        "--test-group-prescheduled",
+        dest="prescheduled",
+        type=str,
+        default=None,
+        help="Path to JSON file containing which tests to run.",
+    )
 
 
 @pytest.hookimpl(hookwrapper=True)
 def pytest_collection_modifyitems(session, config, items):
     yield
-    group_count = config.getoption('test-group-count')
-    group_id = config.getoption('test-group')
-    seed = config.getoption('random-seed')
-    prescheduled_path = config.getoption('prescheduled')
+    group_count = config.getoption("test-group-count")
+    group_id = config.getoption("test-group")
+    seed = config.getoption("random-seed")
+    prescheduled_path = config.getoption("prescheduled")
 
     if not group_count or not group_id:
         return
@@ -54,49 +76,50 @@ def pytest_collection_modifyitems(session, config, items):
     prescheduled_data = [[] for _ in range(group_count)]
     if prescheduled_path:
         try:
-            with open(prescheduled_path, 'r') as f:
+            with open(prescheduled_path, "r") as f:
                 prescheduled_data = json.load(f)
                 if len(prescheduled_data) != group_count:
-                    print('WARNING: Prescheduled tests do not match up with the group count. '
-                          'Prescheduling will be skipped.')
+                    print(
+                        "WARNING: Prescheduled tests do not match up with the group count. "
+                        "Prescheduling will be skipped."
+                    )
         except Exception:
-            print('WARNING: Unable to load prescheduled tests. Prescheduling will be skipped.')
-
-    all_prescheduled_tests = [test_dict[test_name]
-                              for sublist in prescheduled_data
-                              for test_name in sublist
-                              if test_name in test_dict]
-    prescheduled_tests = [test_dict[test_name]
-                          for test_name in prescheduled_data[group_id - 1]
-                          if test_name in test_dict]
+            print(
+                "WARNING: Unable to load prescheduled tests. Prescheduling will be skipped."
+            )
+
+    all_prescheduled_tests = [
+        test_dict[test_name]
+        for sublist in prescheduled_data
+        for test_name in sublist
+        if test_name in test_dict
+    ]
+    prescheduled_tests = [
+        test_dict[test_name]
+        for test_name in prescheduled_data[group_id - 1]
+        if test_name in test_dict
+    ]
     unscheduled_tests = [item for item in items if item not in all_prescheduled_tests]
     unscheduled_tests.sort(key=lambda x: x.name)
 
     if seed is not False:
         seeded = Random(seed)
         seeded.shuffle(unscheduled_tests)
 
-    total_unscheduled_items = len(unscheduled_tests)
-
-    group_size = get_group_size(total_unscheduled_items, group_count)
-    tests_in_group = get_group(unscheduled_tests, group_size, group_id)
+    tests_in_group = get_group(unscheduled_tests, group_count, group_id)
     items[:] = tests_in_group + prescheduled_tests
 
     items.sort(key=original_order.__getitem__)
 
-    terminal_reporter = config.pluginmanager.get_plugin('terminalreporter')
+    terminal_reporter = config.pluginmanager.get_plugin("terminalreporter")
     if terminal_reporter is not None:
         terminal_writer = create_terminal_writer(config)
         message = terminal_writer.markup(
-            '\nRunning test group #{0} ({1} tests)\n'.format(
-                group_id,
-                len(items)
-            ),
-            yellow=True
+            "\nRunning test group #{0} ({1} tests)\n".format(group_id, len(items)),
+            yellow=True,
         )
         terminal_reporter.write(message)
         message = terminal_writer.markup(
-            '\n'.join([item.name for item in items])+'\n',
-            yellow=True
+            "\n".join([item.name for item in items]) + "\n", yellow=True
         )
         terminal_reporter.write(message)

setup.cfg

@@ -1,5 +1,6 @@
 [flake8]
 max-line-length = 119
+exclude = venv .tox
 
 [wheel]
 universal = true

setup.py

@@ -6,35 +6,38 @@
 
 def read(fname):
     file_path = os.path.join(os.path.dirname(__file__), fname)
-    return codecs.open(file_path, encoding='utf-8').read()
+    return codecs.open(file_path, encoding="utf-8").read()
 
 
 setup(
-    name="pytest-split-tests",
-    description=('A Pytest plugin for running a subset of your tests by '
-                 'splitting them in to equally sized groups. Forked from '
-                 'Mark Adams\' original project pytest-test-groups.'),
-    url='https://github.com/wchill/pytest-split-tests',
-    author='Eric Ahn',
-    author_email='[email protected]',
-    packages=['pytest_split_tests'],
-    version='1.1.0',
-    long_description=read('README.rst'),
-    install_requires=['pytest>=2.5'],
-    classifiers=['Development Status :: 5 - Production/Stable',
-                 'Intended Audience :: Developers',
-                 'License :: OSI Approved :: MIT License',
-                 'Operating System :: OS Independent',
-                 'Programming Language :: Python',
-                 'Topic :: Software Development :: Testing',
-                 'Programming Language :: Python :: 3.6',
-                 'Programming Language :: Python :: 3.7',
-                 'Programming Language :: Python :: 3.8',
-                 'Programming Language :: Python :: 3.9'
-                 ],
+    name="picogrid-pytest-split-tests",
+    description=(
+        "A Picogrid Fork of Pytest plugin for running a subset of your tests by "
+        "splitting them in to equally sized groups. Forked from "
+        "Mark Adams' original project pytest-test-groups."
+    ),
+    url="https://github.com/picogrid/pytest-split-tests",
+    author="Picogrid",
+    author_email="[email protected]",
+    packages=["pytest_split_tests"],
+    version="2.0.0",
+    long_description=read("README.rst"),
+    install_requires=["pytest>=2.5"],
+    classifiers=[
+        "Development Status :: 5 - Production/Stable",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+        "Programming Language :: Python",
+        "Topic :: Software Development :: Testing",
+        "Programming Language :: Python :: 3.6",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+    ],
     entry_points={
-        'pytest11': [
-            'split-tests = pytest_split_tests',
+        "pytest11": [
+            "split-tests = pytest_split_tests",
         ]
     },
 )

tests/test_groups.py

@@ -1,33 +1,37 @@
 import pytest
 
-from pytest_split_tests import get_group, get_group_size
+from pytest_split_tests import get_group
 
 
-def test_group_size_computed_correctly_for_even_group():
-    expected = 8
-    actual = get_group_size(32, 4)  # 32 total tests; 4 groups
+def test_group_is_the_proper_size_no_remainder():
+    num_items = 32
+    num_groups = 8
+    expected_items_per_group = 4
+    items = [str(i) for i in range(num_items)]
 
-    assert expected == actual
+    for i in range(1, num_groups + 1):
+        assert len(get_group(items, num_groups, i)) == expected_items_per_group
 
 
-def test_group_size_computed_correctly_for_odd_group():
-    expected = 8
-    actual = get_group_size(31, 4)  # 31 total tests; 4 groups
+def test_group_is_the_proper_size_with_remainder():
+    num_items = 32
+    num_groups = 5
+    expected_base_items_per_group = 6
+    items = [str(i) for i in range(num_items)]
 
-    assert expected == actual
+    for i in range(1, 3):
+        assert len(get_group(items, num_groups, i)) == (
+            expected_base_items_per_group + 1
+        )
 
-
-def test_group_is_the_proper_size():
-    items = [str(i) for i in range(32)]
-    group = get_group(items, 8, 1)
-
-    assert len(group) == 8
+    for i in range(3, num_groups + 1):
+        assert len(get_group(items, num_groups, i)) == expected_base_items_per_group
 
 
 def test_all_groups_together_form_original_set_of_tests():
     items = [str(i) for i in range(32)]
 
-    groups = [get_group(items, 8, i) for i in range(1, 5)]
+    groups = [get_group(items, 4, i) for i in range(1, 5)]
 
     combined = []
     for group in groups:
@@ -40,11 +44,11 @@ def test_group_that_is_too_high_raises_value_error():
     items = [str(i) for i in range(32)]
 
     with pytest.raises(ValueError):
-        get_group(items, 8, 5)
+        get_group(items, 4, 5)
 
 
 def test_group_that_is_too_low_raises_value_error():
     items = [str(i) for i in range(32)]
 
     with pytest.raises(ValueError):
-        get_group(items, 8, 0)
+        get_group(items, 4, 0)

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = py{36,37,38,39}, flake8
+envlist = py{311}, flake8
 
 [testenv]
 commands =