From a960a18614f29a2b5652a677afae98ef9e6b3419 Mon Sep 17 00:00:00 2001
From: Samet Akcay <samet.akcay@intel.com>
Date: Thu, 12 Dec 2024 08:46:44 +0000
Subject: [PATCH] Add cpu and gpu markers to tests to be able to explicitly
 choose them

Signed-off-by: Samet Akcay <samet.akcay@intel.com>
---
 .github/actions/pytest/action.yaml | 6 +++---
 pyproject.toml                     | 4 ++++
 tests/conftest.py                  | 7 +++++++
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/.github/actions/pytest/action.yaml b/.github/actions/pytest/action.yaml
index fa63309760..68df15302b 100644
--- a/.github/actions/pytest/action.yaml
+++ b/.github/actions/pytest/action.yaml
@@ -137,9 +137,9 @@ runs:
 
         # Set device-specific pytest arguments
         if [ "${{ inputs.device }}" = "cpu" ]; then
-          DEVICE_ARGS="--markers='not gpu'"
+          DEVICE_ARGS="-m cpu"
         else
-          DEVICE_ARGS=""
+          DEVICE_ARGS="-m 'cpu or gpu'"  # Run all tests on GPU
         fi
 
         # Run pytest
@@ -150,7 +150,7 @@ runs:
           --timeout=${{ inputs.max-test-time }} \
           --verbosity=0 \
           --durations-only \
-          $DEVICE_ARGS
+          ${DEVICE_ARGS}
 
         # Store test result and duration
         exit_code=$?
diff --git a/pyproject.toml b/pyproject.toml
index 325c60cac9..efdad6e41c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -291,6 +291,10 @@ skips = ["B101"]
 addopts = ["--strict-markers", "--strict-config", "--showlocals", "-ra"]
 testpaths = "tests"
 pythonpath = "src"
+markers = [
+    "gpu: marks tests that require GPU",
+    "cpu: marks tests that can run on CPU only (default)",
+]
 
 
 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
diff --git a/tests/conftest.py b/tests/conftest.py
index a9db6c1d3d..c2709ad275 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -101,3 +101,10 @@ def checkpoint(model_name: str) -> Path:
         return _ckpt_path
 
     return checkpoint
+
+
+def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
+    """Automatically mark tests as 'cpu' unless they're marked as 'gpu'."""
+    for item in items:
+        if not any(marker.name == "gpu" for marker in item.iter_markers()):
+            item.add_marker(pytest.mark.cpu)