test(problems): limit number of iterations in optimize

The number can be overridden by using the environment variable ENGIBENCH_MAX_ITER

Author: g-braeunlich

tests/test_problem_implementations.py

@@ -2,8 +2,9 @@
 
 import dataclasses
 import inspect
+import os
 import sys
-from typing import get_args, get_origin
+from typing import Any, get_args, get_origin
 
 import gymnasium
 from gymnasium import spaces
@@ -88,6 +89,11 @@ def test_problem_impl(problem_class: type[Problem]) -> None:
     print(f"Done testing {problem_class.__name__}.")
 
 
+def problem_id(problem_class: type[Problem]) -> str:
+    id_, _ = problem_class.__module__.removeprefix("engibench.problems.").split(".", 1)
+    return id_
+
+
 @pytest.mark.parametrize(
     "problem_class",
     PYTHON_PROBLEMS + (CONTAINER_PROBLEMS if sys.platform.startswith("linux") else []),
@@ -121,15 +127,16 @@ def test_python_problem_impl(problem_class: type[Problem]) -> None:
     # Test optimization outputs
     print(f"Optimizing {problem_class.__name__}...")
     # Skip optimization test for power electronics, airfoil, and heat conduction problems
-    if (
-        problem_class.__module__.startswith("engibench.problems.power_electronics")
-        or problem_class.__module__.startswith("engibench.problems.airfoil")
-        or problem_class.__module__.startswith("engibench.problems.heatconduction")
-    ):
+    if problem_id(problem_class) in {"power_electronics", "airfoil", "heatconduction"}:
         print(f"Skipping optimization test for {problem_class.__name__}")
         return
     problem.reset(seed=1)
-    optimal_design, history = problem.optimize(starting_point=design)
+    default_max_iter = 20
+    max_iter = os.environ.get("ENGIBENCH_MAX_ITER", default_max_iter)
+    max_iter_config: dict[str, Any] = {
+        key: max_iter for key in ("max_iter", "num_optimization_steps") if hasattr(problem_class.Config, key)
+    }
+    optimal_design, history = problem.optimize(starting_point=design, config=max_iter_config)
     if isinstance(problem.design_space, spaces.Box):
         assert np.all(optimal_design >= problem.design_space.low), (
             f"Problem {problem_class.__name__}: The optimal design should be within the design space."