1 files changed, 35 insertions, 0 deletions
diff --git a/tests/xargs_test.py b/tests/xargs_test.py
index b0a8e0d..e8000b2 100644
--- a/tests/xargs_test.py
+++ b/tests/xargs_test.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import concurrent.futures
+import multiprocessing
 import os
 import sys
 import time
@@ -12,6 +13,40 @@ from pre_commit import parse_shebang
 from pre_commit import xargs
 
 
+def test_cpu_count_sched_getaffinity_exists():
+    with mock.patch.object(
+            os, 'sched_getaffinity', create=True, return_value=set(range(345)),
+    ):
+        assert xargs.cpu_count() == 345
+
+
+@pytest.fixture
+def no_sched_getaffinity():
+    # Simulates an OS without os.sched_getaffinity available (mac/windows)
+    # https://docs.python.org/3/library/os.html#interface-to-the-scheduler
+    with mock.patch.object(
+            os,
+            'sched_getaffinity',
+            create=True,
+            side_effect=AttributeError,
+    ):
+        yield
+
+
+def test_cpu_count_multiprocessing_cpu_count_implemented(no_sched_getaffinity):
+    with mock.patch.object(multiprocessing, 'cpu_count', return_value=123):
+        assert xargs.cpu_count() == 123
+
+
+def test_cpu_count_multiprocessing_cpu_count_not_implemented(
+        no_sched_getaffinity,
+):
+    with mock.patch.object(
+            multiprocessing, 'cpu_count', side_effect=NotImplementedError,
+    ):
+        assert xargs.cpu_count() == 1
+
+
 @pytest.mark.parametrize(
     ('env', 'expected'),
     (