From 331395d85b8afbbca9fb92e1df19564c1988e023 Mon Sep 17 00:00:00 2001
From: rasswanth-s <43314053+rasswanth-s@users.noreply.github.com>
Date: Thu, 8 Feb 2024 13:59:13 +0530
Subject: [PATCH] Added pytest hook to pass cpu count by os.cpu_count()

---
 packages/syft/tests/conftest.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/packages/syft/tests/conftest.py b/packages/syft/tests/conftest.py
index 0ce969a38ea..9d9dd60e0d6 100644
--- a/packages/syft/tests/conftest.py
+++ b/packages/syft/tests/conftest.py
@@ -1,5 +1,6 @@
 # stdlib
 import json
+import os
 from pathlib import Path
 from unittest import mock
 
@@ -46,6 +47,14 @@ def remove_file(filepath: Path):
     filepath.unlink(missing_ok=True)
 
 
+# Pytest hook to set the number of workers for xdist
+def pytest_xdist_auto_num_workers(config):
+    num = config.option.numprocesses
+    if num == "auto" or num == "logical":
+        return os.cpu_count()
+    return None
+
+
 @pytest.fixture(autouse=True)
 def protocol_file():
     random_name = sy.UID().to_string()