olayiwola-gbenga commented on code in PR #74664:
URL: https://github.com/apache/airflow/pull/74664#discussion_r4527207163
##########
providers/google/tests/unit/google/firebase/hooks/test_firestore.py:
##########
@@ -131,6 +136,229 @@ def test_error_operation(self, _, get_conn_mock):
with pytest.raises(AirflowException, match="error"):
self.hook.export_documents(body=EXPORT_DOCUMENT_BODY,
project_id=TEST_PROJECT_ID)
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn")
+ def test_list_operations(self, get_conn_mock):
+ service_mock = get_conn_mock.return_value
+ mock_list =
service_mock.projects.return_value.databases.return_value.operations.return_value.list
+ second_operation = {"name": "operation-name-2"}
+ mock_list.return_value.execute.side_effect = [
+ {"operations": [TEST_OPERATION], "nextPageToken": "page-2"},
+ {"operations": [second_operation]},
+ ]
+
+ ops = self.hook.list_operations(project_id=TEST_PROJECT_ID)
+
+ assert ops == [TEST_OPERATION, second_operation]
+ assert mock_list.call_args_list == [
+
mock.call(name="projects/firestore--project-id/databases/(default)"),
+
mock.call(name="projects/firestore--project-id/databases/(default)",
pageToken="page-2"),
+ ]
+
+ @pytest.mark.parametrize(
+ ("body", "operations", "include_completed", "expected_name"),
+ [
+ pytest.param(
+ EXPORT_DOCUMENT_BODY,
+ [
+ "invalid-entry",
+ {
+ "name": "op-other-type",
+ "metadata": {"@type":
"type.googleapis.com/ImportDocumentsMetadata"},
+ },
+ {
+ "name": "op-failed",
+ "metadata": {
+ "@type":
"type.googleapis.com/ExportDocumentsMetadata",
+ "operationState": "FAILED",
+ "outputUriPrefix":
"gs://test-bucket/test-namespace",
+ "collectionIds": ["test-collection"],
+ },
+ },
+ {
+ "name": "op-wrong-collection",
+ "metadata": {
+ "outputUriPrefix":
"gs://test-bucket/test-namespace/",
+ "collectionIds": ["other"],
+ },
+ },
+ {
+ "name": "op-wrong-namespace",
+ "metadata": {
+ "outputUriPrefix":
"gs://test-bucket/test-namespace/",
+ "collectionIds": ["test-collection"],
+ "namespaceIds": ["ns1"],
+ },
+ },
+ {
+ "name": "op-wrong-uri",
+ "metadata": {
+ "outputUriPrefix": "gs://other-bucket/path",
+ "collectionIds": ["test-collection"],
+ },
+ },
+ {
+ "name": "op-completed",
+ "done": True,
+ "metadata": "non-dict-metadata",
+ "response": {"outputUriPrefix":
"gs://test-bucket/test-namespace"},
+ },
+ {
+ "name": "op-in-progress",
+ "done": False,
+ "metadata": {
+ "@type":
"type.googleapis.com/google.firestore.admin.v1.ExportDocumentsMetadata",
+ "operationState": "PROCESSING",
+ "outputUriPrefix":
"gs://test-bucket/test-namespace",
+ "collectionIds": ["test-collection"],
+ },
+ },
+ ],
+ True,
+ "op-in-progress",
+ id="prefers_in_progress_match",
+ ),
+ pytest.param(
+ EXPORT_DOCUMENT_BODY,
+ [
+ {
+ "name": "op-completed",
+ "done": True,
+ "metadata": {
+ "collectionIds": ["test-collection"],
+ },
+ "response": {"outputUriPrefix":
"gs://test-bucket/test-namespace"},
+ }
+ ],
+ True,
+ "op-completed",
+ id="matches_completed_operation_via_response_uri",
+ ),
+ pytest.param(
+ EXPORT_DOCUMENT_BODY,
+ [
+ {
+ "name": "op-completed",
+ "done": True,
+ "metadata": {
+ "collectionIds": ["test-collection"],
+ },
+ "response": {"outputUriPrefix":
"gs://test-bucket/test-namespace"},
+ }
+ ],
+ False,
+ None,
+ id="ignores_completed_when_include_completed_false",
+ ),
+ pytest.param(
+ {"collectionIds": ["test-collection"]},
+ [
+ {
+ "name": "op-done-no-target-uri",
+ "done": True,
+ "metadata": {"collectionIds": ["test-collection"]},
+ }
+ ],
+ True,
+ None,
+ id="returns_none_when_no_target_uri_specified",
+ ),
+ ],
+ )
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.list_operations")
+ def test_find_matching_export_operation(
+ self, mock_list_ops, body, operations, include_completed, expected_name
+ ):
+ mock_list_ops.return_value = operations
+ result = self.hook.find_matching_export_operation(
+ body=body, project_id=TEST_PROJECT_ID,
include_completed=include_completed
+ )
+ if expected_name is None:
+ assert result is None
+ else:
+ assert result is not None
+ assert result["name"] == expected_name
+
+ @mock.patch(
+
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.find_matching_export_operation"
+ )
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn")
+ def test_export_documents_reattaches_to_in_progress_operation(self,
get_conn_mock, mock_find_op):
+ 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_find_op.return_value = {"name": "existing-op", "done": False}
+ 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_not_called()
+ mock_operation_get.assert_called_once_with(name="existing-op")
+
+ @mock.patch(
+
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.find_matching_export_operation"
+ )
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn")
+ def test_export_documents_skips_when_already_completed_on_http_400(self,
get_conn_mock, mock_find_op):
+ 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
+ )
+ http_400 = HttpError(httplib2.Response({"status": 400}), b"Path
already exists")
+ mock_export_documents.return_value.execute.side_effect = http_400
+ mock_find_op.side_effect = [
+ None,
+ {
+ "name": "completed-op",
+ "done": True,
+ "response": {"outputUriPrefix": "x"},
+ },
+ ]
+
+ self.hook.export_documents(body=EXPORT_DOCUMENT_BODY,
project_id=TEST_PROJECT_ID)
+
+ mock_export_documents.assert_called_once()
+ mock_operation_get.assert_not_called()
+
+ @mock.patch(
+
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.find_matching_export_operation"
+ )
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn")
+ def
test_start_export_documents_recovers_on_http_400_race_in_progress(self,
get_conn_mock, mock_find_op):
+ service_mock = get_conn_mock.return_value
+ mock_export_documents =
service_mock.projects.return_value.databases.return_value.exportDocuments
+ http_400 = HttpError(httplib2.Response({"status": 400}), b"Operation
already in progress")
+ mock_export_documents.return_value.execute.side_effect = http_400
+ mock_find_op.side_effect = [None, {"name": "race-op", "done": False}]
+
+ op = self.hook.start_export_documents(body=EXPORT_DOCUMENT_BODY,
project_id=TEST_PROJECT_ID)
+
+ assert op == {"name": "race-op", "done": False}
+
+ @mock.patch(
+
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.find_matching_export_operation"
+ )
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn")
+ def test_start_export_documents_reraises_http_error_when_no_match(self,
get_conn_mock, mock_find_op):
+ service_mock = get_conn_mock.return_value
+ mock_export_documents =
service_mock.projects.return_value.databases.return_value.exportDocuments
+ http_400 = HttpError(httplib2.Response({"status": 400}), b"Invalid
argument")
+ mock_export_documents.return_value.execute.side_effect = http_400
+ mock_find_op.return_value = None
+
+ with pytest.raises(HttpError):
+ self.hook.start_export_documents(
+ body=EXPORT_DOCUMENT_BODY, project_id=TEST_PROJECT_ID,
reattach_on_restart=False
Review Comment:
I removed `reattach_on_restart=False` from
`test_start_export_documents_reraises_http_error_when_no_match` (asserting
mock_find_op.call_count == 2) and added
`test_start_export_documents_reattach_disabled` to assert
`mock_find_op.assert_not_called()`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
