def execute(self, context):
hook = CloudFirestoreHook(
gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version,
impersonation_chain=self.impersonation_chain
)
return hook.export_documents(database_id=self.database_id, body=self.body, project_id=self.project_id)
class TestCloudFirestoreHookWithoutProjectId(unittest.TestCase):
hook = None # type: Optional[CloudFirestoreHook]
def setUp(self):
with mock.patch(
"airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__",
new=mock_base_gcp_hook_no_default_project_id,
):
self.hook = CloudFirestoreHook(gcp_conn_id="test")
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=None,
)
@mock.patch(
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn"
)
def test_create_build(self, mock_get_conn, mock_project_id):
with self.assertRaises(AirflowException) as e:
self.hook.export_documents(body={})
self.assertEqual(
"The project id must be passed either as keyword project_id parameter or as project_id extra in "
"Google Cloud connection definition. Both are not set!",
str(e.exception),
)
def setUp(self):
with mock.patch(
"airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__",
new=mock_base_gcp_hook_default_project_id,
):
self.hook = CloudFirestoreHook(gcp_conn_id="test")
class TestCloudFirestoreHookWithPassedProjectId(unittest.TestCase):
hook = None # type: Optional[CloudFirestoreHook]
def setUp(self):
with mock.patch(
"airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__",
new=mock_base_gcp_hook_default_project_id,
):
self.hook = CloudFirestoreHook(gcp_conn_id="test")
@mock.patch(
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook._authorize"
)
@mock.patch("airflow.providers.google.firebase.hooks.firestore.build")
@mock.patch(
"airflow.providers.google.firebase.hooks.firestore.build_from_document"
)
def test_client_creation(self, mock_build_from_document, mock_build,
mock_authorize):
result = self.hook.get_conn()
mock_build.assert_called_once_with('firestore',
'v1',
cache_discovery=False)
mock_build_from_document.assert_called_once_with(
mock_build.return_value._rootDesc,
http=mock_authorize.return_value)
self.assertEqual(mock_build_from_document.return_value, result)
self.assertEqual(self.hook._conn, result)
@mock.patch(
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn"
)
def test_immediately_complete(self, get_conn_mock):
service_mock = get_conn_mock.return_value
mock_export_documents = service_mock.projects.return_value.databases.return_value.exportDocuments
mock_operation_get = (service_mock.projects.return_value.databases.
return_value.operations.return_value.get)
(mock_export_documents.return_value.execute.return_value
) = TEST_OPERATION
(mock_operation_get.return_value.execute.return_value
) = TEST_DONE_OPERATION
self.hook.export_documents(body=EXPORT_DOCUMENT_BODY,
project_id=TEST_PROJECT_ID)
mock_export_documents.assert_called_once_with(
body=EXPORT_DOCUMENT_BODY,
name='projects/firestore--project-id/databases/(default)')
@mock.patch(
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn"
)
@mock.patch("airflow.providers.google.firebase.hooks.firestore.time.sleep")
def test_waiting_operation(self, _, get_conn_mock):
service_mock = get_conn_mock.return_value
mock_export_documents = service_mock.projects.return_value.databases.return_value.exportDocuments
mock_operation_get = (service_mock.projects.return_value.databases.
return_value.operations.return_value.get)
(mock_export_documents.return_value.execute.return_value
) = TEST_OPERATION
execute_mock = mock.Mock(
**{
"side_effect": [
TEST_WAITING_OPERATION, TEST_DONE_OPERATION,
TEST_DONE_OPERATION
]
})
mock_operation_get.return_value.execute = execute_mock
self.hook.export_documents(body=EXPORT_DOCUMENT_BODY,
project_id=TEST_PROJECT_ID)
mock_export_documents.assert_called_once_with(
body=EXPORT_DOCUMENT_BODY,
name='projects/firestore--project-id/databases/(default)')
@mock.patch(
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn"
)
@mock.patch("airflow.providers.google.firebase.hooks.firestore.time.sleep")
def test_error_operation(self, _, get_conn_mock):
service_mock = get_conn_mock.return_value
mock_export_documents = service_mock.projects.return_value.databases.return_value.exportDocuments
mock_operation_get = (service_mock.projects.return_value.databases.
return_value.operations.return_value.get)
(mock_export_documents.return_value.execute.return_value
) = TEST_OPERATION
execute_mock = mock.Mock(
**{"side_effect": [TEST_WAITING_OPERATION, TEST_ERROR_OPERATION]})
mock_operation_get.return_value.execute = execute_mock
with self.assertRaisesRegex(AirflowException, "error"):
self.hook.export_documents(body=EXPORT_DOCUMENT_BODY,
project_id=TEST_PROJECT_ID)
