|  | 
|  | 1 | +import importlib | 
|  | 2 | +import os | 
|  | 3 | +import sys | 
|  | 4 | +import time | 
|  | 5 | +import pytest | 
|  | 6 | +from mesa.visualization import SolaraViz | 
|  | 7 | + | 
|  | 8 | +def get_viz_files(directory): | 
|  | 9 | +    viz_files = [] | 
|  | 10 | +    for root, dirs, files in os.walk(directory): | 
|  | 11 | +        for file in files: | 
|  | 12 | +            if file in ['app.py', 'viz.py']: | 
|  | 13 | +                module_name = os.path.relpath(os.path.join(root, file[:-3])).replace( | 
|  | 14 | +                    os.sep, "." | 
|  | 15 | +                ) | 
|  | 16 | +                viz_files.append(module_name) | 
|  | 17 | +    return viz_files | 
|  | 18 | + | 
|  | 19 | +@pytest.mark.parametrize("module_name", get_viz_files("examples")) | 
|  | 20 | +def test_solara_viz(module_name): | 
|  | 21 | +    # Add the 'examples' directory to the Python path | 
|  | 22 | +    examples_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'examples')) | 
|  | 23 | +    sys.path.insert(0, examples_dir) | 
|  | 24 | + | 
|  | 25 | +    # Add the parent directory of the module to the Python path | 
|  | 26 | +    module_parent_dir = os.path.abspath(os.path.join(examples_dir, os.path.dirname(module_name.replace('.', os.sep)))) | 
|  | 27 | +    if module_parent_dir not in sys.path: | 
|  | 28 | +        sys.path.insert(0, module_parent_dir) | 
|  | 29 | + | 
|  | 30 | +    # Import the visualization module | 
|  | 31 | +    module = importlib.import_module(module_name) | 
|  | 32 | + | 
|  | 33 | +    # Find the SolaraViz instance | 
|  | 34 | +    solara_viz = None | 
|  | 35 | +    for item_name in dir(module): | 
|  | 36 | +        item = getattr(module, item_name) | 
|  | 37 | +        if isinstance(item, SolaraViz): | 
|  | 38 | +            solara_viz = item | 
|  | 39 | +            break | 
|  | 40 | + | 
|  | 41 | +    assert solara_viz is not None, f"No SolaraViz instance found in {module_name}" | 
|  | 42 | + | 
|  | 43 | +    # Get the model instance | 
|  | 44 | +    model = solara_viz.model | 
|  | 45 | + | 
|  | 46 | +    # Store the initial state | 
|  | 47 | +    initial_step = model.steps | 
|  | 48 | + | 
|  | 49 | +    # Do 3 steps | 
|  | 50 | +    for _ in range(3): | 
|  | 51 | +        model.step() | 
|  | 52 | +    assert model.steps == initial_step + 3 | 
|  | 53 | + | 
|  | 54 | +    # Pause (no actual pause needed in this context) | 
|  | 55 | + | 
|  | 56 | +    # 3 more steps | 
|  | 57 | +    for _ in range(3): | 
|  | 58 | +        model.step() | 
|  | 59 | +    assert model.steps == initial_step + 6 | 
|  | 60 | + | 
|  | 61 | +    # Reset | 
|  | 62 | +    model_params = solara_viz.model_params | 
|  | 63 | +    new_model = model.__class__(**model_params) | 
|  | 64 | +    solara_viz.model = new_model | 
|  | 65 | +    model = new_model | 
|  | 66 | + | 
|  | 67 | +    assert model.steps == 0 or model.steps == initial_step | 
|  | 68 | + | 
|  | 69 | +    # Run steps for 5 seconds | 
|  | 70 | +    start_time = time.time() | 
|  | 71 | +    while time.time() - start_time < 5: | 
|  | 72 | +        model.step() | 
|  | 73 | + | 
|  | 74 | +    # Pause (no actual pause needed in this context) | 
|  | 75 | + | 
|  | 76 | +    # Check if the step counter has increased | 
|  | 77 | +    assert model.steps > 0, "Model did not advance steps" | 
|  | 78 | + | 
|  | 79 | +    print(f"Test completed for {module_name}. Final step count: {model.steps}") | 
|  | 80 | + | 
|  | 81 | +    # Remove the added directories from the Python path | 
|  | 82 | +    sys.path.pop(0) | 
|  | 83 | +    if module_parent_dir in sys.path: | 
|  | 84 | +        sys.path.remove(module_parent_dir) | 
|  | 85 | + | 
|  | 86 | +# Run the tests | 
|  | 87 | +if __name__ == "__main__": | 
|  | 88 | +    pytest.main([__file__, "-v"]) | 
0 commit comments