From f7dfab57fdb9e943772c4db963e6607d99a97b35 Mon Sep 17 00:00:00 2001
From: Abyss Watcher <abyss-w4tcher@proton.me>
Date: Tue, 4 Mar 2025 16:46:03 +0100
Subject: [PATCH] adjust helpers and add match_output_row

---
 test/test_volatility.py | 50 ++++++++++++++++++++++++++++++-----------
 1 file changed, 37 insertions(+), 13 deletions(-)

diff --git a/test/test_volatility.py b/test/test_volatility.py
index 07cab2c95..9cc9c3304 100644
--- a/test/test_volatility.py
+++ b/test/test_volatility.py
@@ -6,25 +6,24 @@
 #
 
 import os
-import re
 import subprocess
 import sys
-import shutil
 import tempfile
-import hashlib
-import json
 import contextlib
+import functools
+from typing import List, Tuple
 
 #
 # HELPER FUNCTIONS
 #
 
 
+@functools.lru_cache
 def runvol(args, volatility, python):
     volpy = volatility
     python_cmd = python
 
-    cmd = [python_cmd, volpy] + args
+    cmd = (python_cmd, volpy) + args
     print(" ".join(cmd))
     p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     stdout, stderr = p.communicate()
@@ -38,17 +37,18 @@ def runvol(args, volatility, python):
     return p.returncode, stdout, stderr
 
 
-def runvol_plugin(plugin, img, volatility, python, pluginargs=None, globalargs=None):
-    pluginargs = pluginargs or []
-    globalargs = globalargs or []
+@functools.lru_cache
+def runvol_plugin(
+    plugin, img, volatility, python, pluginargs: Tuple = (), globalargs: Tuple = ()
+):
     args = (
         globalargs
-        + [
+        + (
             "--single-location",
             img,
             "-q",
             plugin,
-        ]
+        )
         + pluginargs
     )
 
@@ -60,17 +60,41 @@ def runvolshell(img, volshell, python, volshellargs=None, globalargs=None):
     globalargs = globalargs or []
     args = (
         globalargs
-        + [
+        + (
             "--single-location",
             img,
             "-q",
-        ]
+        )
         + volshellargs
     )
 
     return runvol(args, volshell, python)
 
 
+def match_output_row(
+    json_out: List[dict], expected_row: dict, exact_match: bool = False
+):
+    """Search each row of a plugin's JSON output for an expected row. Each row is a dict.
+
+        Args:
+            json_out: The plugin's output in JSON format (typically obtained through -r json and json.loads)
+            expected_row: The expected row to be found in the output
+            exact_match: Whether to require exactly the expected row, no more no less, or to anticipate columns' addition by checking only
+    the expected row keys and values
+    """
+
+    if not exact_match:
+        for row in json_out:
+            if all(item in expected_row.items() for item in row.items()):
+                return True
+    else:
+        for row in json_out:
+            if expected_row == row:
+                return True
+
+    return False
+
+
 #
 # TESTS
 #
@@ -96,7 +120,7 @@ def basic_volshell_test(image, volatility, python, globalargs):
             img=image,
             volshell=volatility,
             python=python,
-            volshellargs=["--script", filename],
+            volshellargs=("--script", filename),
             globalargs=globalargs,
         )
     finally: