shahar1 commented on code in PR #71947:
URL: https://github.com/apache/airflow/pull/71947#discussion_r4306520398
##########
providers/google/src/airflow/providers/google/marketing_platform/hooks/bid_manager.py:
##########
@@ -116,12 +122,17 @@ def get_report(self, query_id: str, report_id: str) ->
dict:
.execute(num_retries=self.num_retries)
)
- def list_reports(self, query_id: str) -> dict:
+ def list_reports(self, query_id: str) -> list[dict]:
Review Comment:
Public hook method: the return type changes from `dict` to `list[dict]`, so
callers doing `hook.list_reports(qid)["reports"]` break on upgrade. Keep the
pagination but `return {"reports": reports}`, or add a `.. warning::`
breaking-change note to `providers/google/docs/changelog.rst`.
##########
providers/google/src/airflow/providers/google/marketing_platform/sensors/bid_manager.py:
##########
@@ -21,6 +21,7 @@
from collections.abc import Sequence
from typing import TYPE_CHECKING
+from airflow.exceptions import AirflowFailException
Review Comment:
Deprecated path on Airflow 3 (emits `DeprecatedImportWarning`). Use `from
airflow.providers.common.compat.sdk import AirflowFailException`. Same in the
sensor test.
##########
providers/google/src/airflow/providers/google/marketing_platform/operators/bid_manager.py:
##########
@@ -317,8 +545,10 @@ def execute(self, context: Context):
resource = hook.get_report(query_id=self.query_id,
report_id=self.report_id)
status = resource.get("metadata", {}).get("status", {}).get("state")
- if resource and status not in ["DONE", "FAILED"]:
- raise AirflowException(f"Report {self.report_id} for query
{self.query_id} is still running")
+ if status != "DONE":
Review Comment:
A `FAILED` report never becomes `DONE`, so retrying it is pointless. Raise
`AirflowFailException` for `FAILED` (as the sensor does) and keep a retryable,
more specific exception for in-progress states.
##########
providers/google/tests/unit/google/marketing_platform/sensors/test_bid_manager.py:
##########
@@ -30,15 +33,36 @@
class TestGoogleBidManagerRunQuerySensor:
@mock.patch(f"{MODULE_NAME}.GoogleBidManagerHook")
- @mock.patch(f"{MODULE_NAME}.BaseSensorOperator")
- def test_poke(self, mock_base_op, hook_mock):
+ def test_poke_done(self, hook_mock):
Review Comment:
Nit: these three tests differ only in input and expected result, so one
`@pytest.mark.parametrize` test would do. Same for the two
`DownloadReportOperator` status tests.
##########
providers/google/tests/unit/google/marketing_platform/hooks/test_bid_manager.py:
##########
@@ -93,14 +93,20 @@ def test_get_query(self, get_conn_mock):
@mock.patch("airflow.providers.google.marketing_platform.hooks.bid_manager.GoogleBidManagerHook.get_conn")
def test_list_queries(self, get_conn_mock):
- queries = ["test"]
- return_value = {"queries": queries}
-
get_conn_mock.return_value.queries.return_value.list.return_value.execute.return_value
= return_value
+ queries_page1 = [{"queryId": "q1"}]
+ queries_page2 = [{"queryId": "q2"}]
+ mock_request = mock.MagicMock()
Review Comment:
Nit: `mock.MagicMock(spec=HttpRequest)` (from `googleapiclient.http`), here
and in `test_list_reports`.
##########
providers/google/tests/unit/google/marketing_platform/operators/test_bid_manager.py:
##########
@@ -426,65 +430,220 @@ def f():
object_name=REPORT_NAME + ".gz",
)
+
@mock.patch("airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerHook")
+ def test_execute_failed_status_raises(self, mock_hook):
+ mock_hook.return_value.get_report.return_value = {
+ "metadata": {
+ "status": {"state": "FAILED"},
+ }
+ }
+ mock_context = {"task_instance": mock.Mock()}
+ op = GoogleBidManagerDownloadReportOperator(
+ query_id=QUERY_ID,
+ report_id=REPORT_ID,
+ bucket_name=BUCKET_NAME,
+ report_name=REPORT_NAME,
+ task_id="test_task",
+ )
+ with pytest.raises(AirflowException, match="is in state 'FAILED'.
Expected 'DONE'"):
+ op.execute(context=mock_context)
+
+
@mock.patch("airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerHook")
+ def test_execute_running_status_raises(self, mock_hook):
+ mock_hook.return_value.get_report.return_value = {
+ "metadata": {
+ "status": {"state": "RUNNING"},
+ }
+ }
+ mock_context = {"task_instance": mock.Mock()}
+ op = GoogleBidManagerDownloadReportOperator(
+ query_id=QUERY_ID,
+ report_id=REPORT_ID,
+ bucket_name=BUCKET_NAME,
+ report_name=REPORT_NAME,
+ task_id="test_task",
+ )
+ with pytest.raises(AirflowException, match="is in state 'RUNNING'.
Expected 'DONE'"):
+ op.execute(context=mock_context)
+
class TestGoogleBidManagerRunQueryOperator:
@mock.patch("airflow.providers.google.marketing_platform.operators.bid_manager.GoogleBidManagerHook")
def test_execute(self, hook_mock):
Review Comment:
This PR doesn't change `RunQueryOperator`, so please revert this rewrite
(same for `TestGoogleBidManagerCreateQueryOperator.test_execute`). Keep the
`test_prepare_template` fix.
--
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]
