class TestCatchHttpException(unittest.TestCase):
def test_no_exception(self):
self.called = False
class FixtureClass(LoggingMixin):
@hook.GoogleCloudBaseHook.catch_http_exception
def text_fixture(*args, **kwargs):
self.called = True
FixtureClass().text_fixture()
self.assertTrue(self.called)
@parameterized.expand([
(MovedPermanently("MESSAGE"), ),
(RetryError("MESSAGE", cause=Exception("MESSAGE")), ),
(ValueError("MESSAGE"), ),
])
def test_raise_airflowexception(self, ex_obj):
self.called = False
class FixtureClass(LoggingMixin):
@hook.GoogleCloudBaseHook.catch_http_exception
def test_fixutre(*args, **kwargs):
self.called = True
raise ex_obj
with self.assertRaises(AirflowException):
FixtureClass().test_fixutre()
self.assertTrue(self.called)
def test_raise_alreadyexists(self):
self.called = False
class FixtureClass(LoggingMixin):
@hook.GoogleCloudBaseHook.catch_http_exception
def test_fixutre(*args, **kwargs):
self.called = True
raise AlreadyExists("MESSAGE")
with self.assertRaises(AlreadyExists):
FixtureClass().test_fixutre()
self.assertTrue(self.called)
def test_raise_http_error(self):
self.called = False
class FixtureClass(BaseHook):
@hook.GoogleCloudBaseHook.catch_http_exception
def test_fixtue(*args, **kwargs):
self.called = True
raise HttpError(mock.Mock(**{"reason.return_value": None}),
b"CONTENT")
with self.assertRaises(AirflowException):
FixtureClass(None).test_fixtue()
self.assertTrue(self.called)
class TestCatchHttpException(unittest.TestCase):
# pylint:disable=no-method-argument,unused-argument
@parameterized.expand([
("no_exception", None, LoggingMixin, None, None),
("raise_airflowexception", MovedPermanently("MESSAGE"), LoggingMixin,
None, AirflowException),
("raise_airflowexception",
RetryError("MESSAGE", cause=Exception("MESSAGE")), LoggingMixin, None,
AirflowException),
("raise_airflowexception", ValueError("MESSAGE"), LoggingMixin, None,
AirflowException),
("raise_alreadyexists", AlreadyExists("MESSAGE"), LoggingMixin, None,
AlreadyExists),
("raise_http_error",
HttpError(mock.Mock(**{"reason.return_value": None}),
b"CONTENT"), BaseHook, {
"source": None
}, AirflowException),
])
def test_catch_exception(self, name, exception, base_class,
base_class_args, assert_raised):
self.called = False # pylint:disable=attribute-defined-outside-init
class FixtureClass(base_class):
@hook.GoogleCloudBaseHook.catch_http_exception
def test_fixture(*args, **kwargs): # pylint:disable=unused-argument,no-method-argument
self.called = True # pylint:disable=attribute-defined-outside-init
if exception is not None:
raise exception
if assert_raised is None:
FixtureClass(base_class_args).test_fixture()
else:
with self.assertRaises(assert_raised):
FixtureClass(base_class_args).test_fixture()
self.assertTrue(self.called)
def test_generate_component_gateway_url_raises_retry_error():
with patch('google.cloud.dataproc_v1beta2.ClusterControllerClient.get_cluster', \
side_effect=RetryError('error message', 'cause')):
google_auth_class.get_component_gateway_url('project', 'region',
'cluster',
make_credentials())