diff --git a/bigquery/tests/unit/helpers.py b/bigquery/tests/unit/helpers.py index 5b731a763a99..673aa8ac5f02 100644 --- a/bigquery/tests/unit/helpers.py +++ b/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): + """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) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 1bfbcee8bc56..f591bd623309 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/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,14 @@ 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): + def fail_if(name, globals, locals, fromlist, level): # NOTE: *very* simplified, assuming a straightforward absolute import - if "bigquery_storage_v1beta1" in name or ( + 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"}