Move context manager to swap the current store into spack.store

The context manager can be used to swap the current
store temporarily, for any use case that may need it.

3 files changed, 32 insertions, 15 deletions

diff --git a/lib/spack/spack/store.py b/lib/spack/spack/store.py
index dfa9300743..f53183d323 100644
--- a/lib/spack/spack/store.py
+++ b/lib/spack/spack/store.py
@@ -23,6 +23,7 @@ install trees to define their own layouts with some per-tree
 configuration.
 
 """
+import contextlib
 import os
 import re
 import six
@@ -227,3 +228,28 @@ def _construct_upstream_dbs_from_install_roots(
         accumulated_upstream_dbs.insert(0, next_db)
 
     return accumulated_upstream_dbs
+
+
+@contextlib.contextmanager
+def use_store(store_or_path):
+    """Use the store passed as argument within the context manager.
+
+    Args:
+        store_or_path: either a Store object ot a path to where the store resides
+
+    Returns:
+        Store object associated with the context manager's store
+    """
+    global store
+
+    # Normalize input arguments
+    temporary_store = store_or_path
+    if not isinstance(store_or_path, Store):
+        temporary_store = Store(store_or_path)
+
+    # Swap the store with the one just constructed and return it
+    original_store, store = store, temporary_store
+    yield temporary_store
+
+    # Restore the original store
+    store = original_store
diff --git a/lib/spack/spack/test/cmd/module.py b/lib/spack/spack/test/cmd/module.py
index 854f8a6765..b328fb6cbf 100644
--- a/lib/spack/spack/test/cmd/module.py
+++ b/lib/spack/spack/test/cmd/module.py
@@ -10,7 +10,8 @@ import pytest
 
 import spack.main
 import spack.modules
-from spack.test.conftest import use_store, use_configuration
+import spack.store
+from spack.test.conftest import use_configuration
 
 module = spack.main.SpackCommand('module')
 
@@ -21,7 +22,7 @@ def ensure_module_files_are_there(
         mock_repo_path, mock_store, mock_configuration):
     """Generate module files for module tests."""
     module = spack.main.SpackCommand('module')
-    with use_store(mock_store):
+    with spack.store.use_store(mock_store):
         with use_configuration(mock_configuration):
             with spack.repo.use_repositories(mock_repo_path):
                 module('tcl', 'refresh', '-y')
diff --git a/lib/spack/spack/test/conftest.py b/lib/spack/spack/test/conftest.py
index 8d2d02f587..319bb233cc 100644
--- a/lib/spack/spack/test/conftest.py
+++ b/lib/spack/spack/test/conftest.py
@@ -35,6 +35,7 @@ import spack.paths
 import spack.platforms.test
 import spack.repo
 import spack.stage
+import spack.store
 import spack.util.executable
 import spack.util.gpg
 import spack.subprocess_context
@@ -373,15 +374,6 @@ def use_configuration(config):
     spack.compilers._cache_config_file = saved_compiler_cache
 
 
-@contextlib.contextmanager
-def use_store(store):
-    """Context manager to swap out the global Spack store."""
-    saved = spack.store.store
-    spack.store.store = store
-    yield
-    spack.store.store = saved
-
-
 #
 # Test-specific fixtures
 #
@@ -630,12 +622,11 @@ def mock_store(tmpdir_factory, mock_repo_path, mock_configuration,
 
     """
     store_path, store_cache = _store_dir_and_cache
-    store = spack.store.Store(str(store_path))
 
     # If the cache does not exist populate the store and create it
     if not os.path.exists(str(store_cache.join('.spack-db'))):
         with use_configuration(mock_configuration):
-            with use_store(store):
+            with spack.store.use_store(str(store_path)) as store:
                 with spack.repo.use_repositories(mock_repo_path):
                     _populate(store.db)
         store_path.copy(store_cache, mode=True, stat=True)
@@ -660,12 +651,11 @@ def mutable_mock_store(tmpdir_factory, mock_repo_path, mock_configuration,
 
     """
     store_path, store_cache = _store_dir_and_cache
-    store = spack.store.Store(str(store_path))
 
     # If the cache does not exist populate the store and create it
     if not os.path.exists(str(store_cache.join('.spack-db'))):
         with use_configuration(mock_configuration):
-            with use_store(store):
+            with spack.store.use_store(str(store_path)) as store:
                 with spack.repo.use_repositories(mock_repo_path):
                     _populate(store.db)
         store_path.copy(store_cache, mode=True, stat=True)