From 4e5e1e8f35c3a01fa530a21a39fa7274b54a3989 Mon Sep 17 00:00:00 2001
From: Massimiliano Culpo <massimiliano.culpo@gmail.com>
Date: Mon, 1 Feb 2021 14:55:45 +0100
Subject: [PATCH] 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)
---
 lib/spack/spack/store.py           | 26 ++++++++++++++++++++++++++
 lib/spack/spack/test/cmd/module.py |  5 +++--
 lib/spack/spack/test/conftest.py   | 16 +++-------------
 3 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/lib/spack/spack/store.py b/lib/spack/spack/store.py
index 779a5531a44..6040145b64b 100644
--- a/lib/spack/spack/store.py
+++ b/lib/spack/spack/store.py
@@ -23,6 +23,7 @@
 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 9a34368fced..a488c47d1c0 100644
--- a/lib/spack/spack/test/cmd/module.py
+++ b/lib/spack/spack/test/cmd/module.py
@@ -10,7 +10,8 @@
 
 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 f214c78c181..153bee128f9 100644
--- a/lib/spack/spack/test/conftest.py
+++ b/lib/spack/spack/test/conftest.py
@@ -35,6 +35,7 @@
 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)