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.

(cherry picked from commit cb2c233a97073f8c5d89581ee2a2401fef5f878d)

3 files changed, 32 insertions, 15 deletions

diff --git a/lib/spack/spack/store.py b/lib/spack/spack/store.py
index 779a5531a4..6040145b64 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 9a34368fce..a488c47d1c 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 f214c78c18..153bee128f 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
 #
@@ -631,12 +623,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)
@@ -661,12 +652,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)