def test_get_staging_location_exception(self, mock_cluster_client):
"""
Test to catch when an error is raised inside get_staging_location.
"""
cluster_metadata = MasterURLIdentifier(project_id='test-project',
region='test-region',
cluster_name='test-cluster')
cluster_manager = DataprocClusterManager(cluster_metadata)
with self.assertRaises(MockException):
cluster_manager.get_staging_location(cluster_metadata)
def test_get_staging_location(self, mock_cluster_client, mock_list):
"""
Test to receive a mock staging location successfully under
get_staging_location.
"""
cluster_metadata = MasterURLIdentifier(project_id='test-project',
region='test-region',
cluster_name='test-cluster')
cluster_manager = DataprocClusterManager(cluster_metadata)
self.assertEqual(
cluster_manager.get_staging_location(cluster_metadata),
'gs://test-bucket/google-cloud-dataproc-metainfo/')
