From cc7bd1b270429fb8fb90d9bf02ef57ce094a394c Mon Sep 17 00:00:00 2001
From: Adam Tyson <code@adamltyson.com>
Date: Fri, 13 Oct 2023 15:02:44 +0100
Subject: [PATCH] Test windows cpu count

---
 brainglobe_utils/general/system.py            |  9 ++++--
 .../test_unit/test_general/test_system.py     | 29 +++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/brainglobe_utils/general/system.py b/brainglobe_utils/general/system.py
index d137cb5..5d0b226 100644
--- a/brainglobe_utils/general/system.py
+++ b/brainglobe_utils/general/system.py
@@ -152,8 +152,7 @@ def get_num_processes(
     else:
         n_processes = n_cpu_cores
 
-    if platform.system() == "Windows":
-        n_max_processes = min(n_max_processes, MAX_PROCESSES_WINDOWS)
+    n_max_processes = limit_cpus_windows(n_max_processes)
 
     if n_max_processes is not None:
         if n_max_processes < n_processes:
@@ -172,6 +171,12 @@ def get_num_processes(
     return int(n_processes)
 
 
+def limit_cpus_windows(n_max_processes):
+    if platform.system() == "Windows":
+        n_max_processes = min(n_max_processes, MAX_PROCESSES_WINDOWS)
+    return n_max_processes
+
+
 def get_cores_available():
     try:
         os.environ["SLURM_JOB_ID"]
diff --git a/tests/tests/test_unit/test_general/test_system.py b/tests/tests/test_unit/test_general/test_system.py
index f50d831..4d7bf23 100644
--- a/tests/tests/test_unit/test_general/test_system.py
+++ b/tests/tests/test_unit/test_general/test_system.py
@@ -2,6 +2,7 @@
 import random
 from pathlib import Path
 from random import shuffle
+from unittest.mock import patch
 
 import pytest
 
@@ -90,6 +91,7 @@ def test_get_num_processes():
     assert os.cpu_count() == system.get_num_processes(min_free_cpu_cores=0)
 
 
+## orig
 def test_max_processes():
     max_proc = 5
     correct_n = min(os.cpu_count(), max_proc)
@@ -98,6 +100,33 @@ def test_max_processes():
     )
 
 
+def test_max_processes_windows_low():
+    cpu_count = 10
+    with patch(
+        "brainglobe_utils.general.system.platform.system",
+        return_value="Windows",
+    ):
+        with patch(
+            "brainglobe_utils.general.system.psutil.cpu_count",
+            return_value=cpu_count,
+        ):
+            assert system.limit_cpus_windows(cpu_count) == cpu_count
+
+
+def test_max_processes_windows_high():
+    cpu_count = 128
+    with patch(
+        "brainglobe_utils.general.system.platform.system",
+        return_value="Windows",
+    ):
+        with patch(
+            "brainglobe_utils.general.system.psutil.cpu_count",
+            return_value=cpu_count,
+        ):
+            # 61 is max on Windows
+            assert system.limit_cpus_windows(cpu_count) == 61
+
+
 class Paths:
     def __init__(self, directory):
         self.one = directory / "one.aaa"