change test order

econpizza/testing/test_rest_config.py

@@ -0,0 +1,103 @@
+"""Tests for the config module. Delete any __econpizza__ or __jax_cache__ folders you might have in the current folder before running"""
+import pytest
+import jax
+from unittest.mock import patch
+import shutil
+import os
+import sys
+# autopep8: off
+sys.path.insert(0, os.path.abspath("."))
+import econpizza as ep
+from econpizza.config import EconPizzaConfig
+# autopep8: on
+
+
+@pytest.fixture(scope="function", autouse=True)
+def ep_config_reset():
+    ep.config = EconPizzaConfig()
+
+
+@pytest.fixture(scope="function", autouse=True)
+def os_getcwd_create():
+    test_cache_folder = os.path.abspath("config_working_dir")
+
+    if not os.path.exists(test_cache_folder):
+        os.makedirs(test_cache_folder)
+
+    with patch("os.getcwd", return_value=test_cache_folder):
+        yield
+
+    if os.path.exists(test_cache_folder):
+        shutil.rmtree(test_cache_folder)
+
+
+def test_config_default_values():
+    assert ep.config["enable_jax_persistent_cache"] == False
+    assert ep.config.jax_cache_folder == "__jax_cache__"
+
+
+def test_config_jax_default_values():
+    assert jax.config.values["jax_compilation_cache_dir"] is None
+    assert jax.config.values["jax_persistent_cache_min_entry_size_bytes"] == .0
+    assert jax.config.values["jax_persistent_cache_min_compile_time_secs"] == 1.0
+
+
+@patch("os.makedirs")
+@patch("jax.config.update")
+def test_config_enable_jax_persistent_cache(mock_jax_update, mock_makedirs):
+    ep.config["enable_jax_persistent_cache"] = True
+    mock_makedirs.assert_any_call(os.path.join(
+        os.getcwd(), "__jax_cache__"), exist_ok=True)
+
+    mock_jax_update.assert_any_call(
+        "jax_compilation_cache_dir", os.path.join(os.getcwd(), "__jax_cache__"))
+    mock_jax_update.assert_any_call(
+        "jax_persistent_cache_min_compile_time_secs", 0)
+
+
+@patch("os.makedirs")
+@patch("jax.config.update")
+def test_config_set_jax_folder(mock_jax_update, mock_makedirs):
+    ep.config.jax_cache_folder = "test1"
+    ep.config["enable_jax_persistent_cache"] = True
+    mock_makedirs.assert_any_call(os.path.join(
+        os.getcwd(), "test1"), exist_ok=True)
+    mock_jax_update.assert_any_call(
+        "jax_compilation_cache_dir", os.path.join(os.getcwd(), "test1"))
+
+
+@patch("jax.config.update")
+def test_config_jax_folder_set_from_outside(mock_jax_update):
+    mock_jax_update("jax_compilation_cache_dir", "jax_from_outside")
+    ep.config["enable_jax_persistent_cache"] = True
+    mock_jax_update.assert_any_call(
+        "jax_compilation_cache_dir", "jax_from_outside")
+
+
+@patch("os.path.exists")
+@patch("os.makedirs")
+@patch("jax.config.update")
+def test_jax_cache_folder_not_created_second_time(mock_jax_update, mock_makedirs, mock_exists):
+    # Set to first return False when the folder is not created, then True when the folder is created
+    mock_exists.side_effect = [False, True]
+
+    # When called for the first time, a cache folder should be created(default is __jax_cache__)
+    ep.config["enable_jax_persistent_cache"] = True
+    mock_makedirs.assert_any_call(os.path.join(
+        os.getcwd(), "__jax_cache__"), exist_ok=True)
+    assert mock_jax_update.call_count == 2
+    # Now reset the mock so that the calls are 0 again.
+    mock_makedirs.reset_mock()
+    mock_jax_update.reset_mock()
+    # The second time we should not create a folder
+    ep.config["enable_jax_persistent_cache"] = True
+    mock_makedirs.assert_not_called()
+    assert mock_jax_update.call_count == 0
+
+
+def test_config_enable_jax_persistent_cache_called_after_model_load():
+    _ = ep.load(ep.examples.dsge)
+
+    assert os.path.exists(ep.config.jax_cache_folder) == False
+    ep.config["enable_jax_persistent_cache"] = True
+    assert os.path.exists(ep.config.jax_cache_folder) == True