Extract import patcher factory to test helpers

bigquery/tests/unit/helpers.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import mock
+import six
+
 
 def make_connection(*responses):
     import google.cloud.bigquery._http
@@ -22,3 +25,25 @@ def make_connection(*responses):
     mock_conn.user_agent = "testing 1.2.3"
     mock_conn.api_request.side_effect = list(responses) + [NotFound("miss")]
     return mock_conn
+
+
+def maybe_fail_import(predicate=lambda: None):
+    """Create and return a patcher that conditionally makes an import fail.
+
+    Args:
+        predicate (Callable[[...], bool]): A callable that, if it returns `True`,
+            triggers an `ImportError`. It must accept the same arguments as the
+            built-in `__import__` function.
+            https://docs.python.org/3/library/functions.html#__import__
+
+    Returns:
+        A mock patcher object that can be used to enable patched import behavior.
+    """
+    orig_import = six.moves.builtins.__import__
+
+    def custom_import(name, globals=None, locals=None, fromlist=(), level=0):
+        if predicate(name, globals, locals, fromlist, level):
+            raise ImportError
+        return orig_import(name, globals, locals, fromlist, level)
+
+    return mock.patch.object(six.moves.builtins, "__import__", new=custom_import)

bigquery/tests/unit/test_magics.py

@@ -42,6 +42,7 @@
 from google.cloud.bigquery import table
 from google.cloud.bigquery import magics
 from tests.unit.helpers import make_connection
+from tests.unit.helpers import maybe_fail_import
 
 
 pytestmark = pytest.mark.skipif(IPython is None, reason="Requires `ipython`")
@@ -68,17 +69,13 @@ def ipython_interactive(request, ipython):
 @pytest.fixture(scope="session")
 def missing_bq_storage():
     """Provide a patcher that can make the bigquery storage import to fail."""
-    orig_import = six.moves.builtins.__import__
 
-    def custom_import(name, globals=None, locals=None, fromlist=(), level=0):
-        # NOTE: *very* simplified, assuming a straightforward absolute import
-        if "bigquery_storage_v1beta1" in name or (
+    def fail_if(name, globals, locals, fromlist, level):
+        return "bigquery_storage_v1beta1" in name or (
             fromlist is not None and "bigquery_storage_v1beta1" in fromlist
-        ):
-            raise ImportError
-        return orig_import(name, globals, locals, fromlist, level)
+        )
 
-    return mock.patch.object(six.moves.builtins, "__import__", new=custom_import)
+    return maybe_fail_import(predicate=fail_if)
 
 
 JOB_REFERENCE_RESOURCE = {"projectId": "its-a-project-eh", "jobId": "some-random-id"}