From 535ce3a22a575c14100ca30627f3d2889fd86c9c Mon Sep 17 00:00:00 2001
From: Gustavo Moreira <gmoreira@gmail.com>
Date: Sun, 1 Dec 2024 19:34:35 +1100
Subject: [PATCH] testing: Enable automatic selection of the OS image based on
 the test and filename prefix, addressing an issue in the development
 environment.

For example, when using VSCode with pytest, test autodiscovery triggers pytest_generate_tests(), adding all images to each test case. This causes issues, as Linux tests end up being executed with Windows and Mac images, and vice versa.
---
 .github/workflows/test.yaml | 10 ++++++----
 test/conftest.py            | 33 ++++++++++++++++++++++++++++-----
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 084cce2956..dfc42499df 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -25,10 +25,13 @@ jobs:
 
     - name: Download images
       run: |
+        mkdir test_images
+        cd test_images
         curl -sLO "https://downloads.volatilityfoundation.org/volatility3/images/linux-sample-1.bin.gz"
         gunzip linux-sample-1.bin.gz
         curl -sLO "https://downloads.volatilityfoundation.org/volatility3/images/win-xp-laptop-2005-06-25.img.gz"
         gunzip win-xp-laptop-2005-06-25.img.gz
+        cd -
 
     - name: Download and Extract symbols
       run: |
@@ -39,13 +42,12 @@ jobs:
 
     - name: Testing...
       run: |
-        pytest ./test/test_volatility.py --volatility=vol.py --image win-xp-laptop-2005-06-25.img -k test_windows -v
-        pytest ./test/test_volatility.py --volatility=vol.py --image linux-sample-1.bin -k test_linux -v
+        pytest ./test/test_volatility.py --volatility=vol.py --image-dir=./test_images -k test_windows -v
+        pytest ./test/test_volatility.py --volatility=vol.py --image-dir=./test_images -k test_linux -v
 
     - name: Clean up post-test
       run: |
-        rm -rf *.bin
-        rm -rf *.img
+        rm -rf test_images
         cd volatility3/symbols
         rm -rf linux
         rm -rf linux.zip
diff --git a/test/conftest.py b/test/conftest.py
index 4ad63065bb..0115fade93 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -35,16 +35,39 @@ def pytest_addoption(parser):
 def pytest_generate_tests(metafunc):
     """Parameterize tests based on image names"""
 
-    images = metafunc.config.getoption("image")
+    images = metafunc.config.getoption("image").copy()
     for image_dir in metafunc.config.getoption("image_dir"):
-        images = images + [
-            os.path.join(image_dir, dir) for dir in os.listdir(image_dir)
+        images += [
+            os.path.join(image_dir, dir_name) for dir_name in os.listdir(image_dir)
         ]
 
-    # tests with "image" parameter are run against images
+    # tests with "image" parameter are run against image
     if "image" in metafunc.fixturenames:
+        filtered_images = []
+        ids = []
+        for image in images:
+            image_base = os.path.basename(image)
+            test_name = metafunc.definition.originalname
+            if test_name.startswith("test_windows_") and not image_base.startswith(
+                "win-"
+            ):
+                continue
+            elif test_name.startswith("test_linux_") and not image_base.startswith(
+                "linux-"
+            ):
+                continue
+            elif test_name.startswith("test_mac_") and not image_base.startswith(
+                "mac-"
+            ):
+                continue
+
+            filtered_images.append(image)
+            ids.append(image_base)
+
         metafunc.parametrize(
-            "image", images, ids=[os.path.basename(image) for image in images]
+            "image",
+            filtered_images,
+            ids=ids,
         )