4242from google .cloud .bigquery import table
4343from google .cloud .bigquery import magics
4444from tests .unit .helpers import make_connection
45+ from tests .unit .helpers import maybe_fail_import
4546
4647
4748pytestmark = pytest .mark .skipif (IPython is None , reason = "Requires `ipython`" )
@@ -68,17 +69,14 @@ def ipython_interactive(request, ipython):
6869@pytest .fixture (scope = "session" )
6970def missing_bq_storage ():
7071 """Provide a patcher that can make the bigquery storage import to fail."""
71- orig_import = six .moves .builtins .__import__
7272
73- def custom_import (name , globals = None , locals = None , fromlist = () , level = 0 ):
73+ def fail_if (name , globals , locals , fromlist , level ):
7474 # NOTE: *very* simplified, assuming a straightforward absolute import
75- if "bigquery_storage_v1beta1" in name or (
75+ return "bigquery_storage_v1beta1" in name or (
7676 fromlist is not None and "bigquery_storage_v1beta1" in fromlist
77- ):
78- raise ImportError
79- return orig_import (name , globals , locals , fromlist , level )
77+ )
8078
81- return mock . patch . object ( six . moves . builtins , "__import__" , new = custom_import )
79+ return maybe_fail_import ( predicate = fail_if )
8280
8381
8482JOB_REFERENCE_RESOURCE = {"projectId" : "its-a-project-eh" , "jobId" : "some-random-id" }
0 commit comments