Allow additional env vars in GPU setting checks

Author: shuds13

libensemble/tests/unit_tests/test_executor_gpus.py

@@ -118,7 +118,9 @@ def run_check(exp_env, exp_cmd, **kwargs):
         args_for_sim = "sleep 0"
         exp_runline = exp_cmd + " simdir/my_simtask.x sleep 0"
         task = exctr.submit(calc_type="sim", app_args=args_for_sim, dry_run=True, **kwargs)
-        assert task.env == exp_env, f"Task env does not match expected:\n Received: {task.env}\n Expected: {exp_env}"
+        for key, value in exp_env.items():
+            assert key in task.env, f"Expected env key '{key}' not found in task.env: {task.env}"
+            assert task.env[key] == value, f"Env key '{key}' has value '{task.env[key]}', expected '{value}'"
         assert (
             task.runline == exp_runline
         ), f"Run line  does not match expected.\n Received: {task.runline}\n Expected: {exp_runline}"

libensemble/tools/test_support.py

@@ -242,6 +242,16 @@ def check_gpu_setting(task, assert_setting=True, print_setting=False, resources=
         print(f"Worker {task.workerID}: {desc}GPU setting ({stype}): {gpu_setting} {addon}", flush=True)
 
     if assert_setting:
-        assert (
-            gpu_setting == expected
-        ), f"Worker {task.workerID}: Found GPU setting: {gpu_setting}, Expected: {expected}"
+        if isinstance(expected, dict):
+            for key, value in expected.items():
+                assert key in gpu_setting, (
+                    f"Worker {task.workerID}: Expected env key '{key}' not found in GPU setting: {gpu_setting}"
+                )
+                assert gpu_setting[key] == value, (
+                    f"Worker {task.workerID}: GPU setting key '{key}' has value '{gpu_setting[key]}', "
+                    f"expected '{value}'"
+                )
+        else:
+            assert (
+                gpu_setting == expected
+            ), f"Worker {task.workerID}: Found GPU setting: {gpu_setting}, Expected: {expected}"