SameerMesiah97 commented on code in PR #61790:
URL: https://github.com/apache/airflow/pull/61790#discussion_r2796254830
##########
providers/google/tests/unit/google/cloud/operators/test_gen_ai.py:
##########
@@ -391,6 +525,35 @@ def test_execute(self, mock_hook):
job_name=TEST_BATCH_JOB_NAME,
)
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def test_execute_propagates_value_error(self, mock_hook):
+ op = GenAIGeminiGetBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ job_name=TEST_BATCH_JOB_NAME,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ )
+
+ mock_hook.return_value.get_batch_job.side_effect = ValueError()
+
+ with pytest.raises(
+ AirflowException, match=f"('Job with name %s not found',
'{TEST_BATCH_JOB_NAME}')"
Review Comment:
This might be a bit too brittle. What if someone refactors the exception
message to a standard f-string like:
`f"Job with name {TEST_BATCH_JOB_NAME} not found"`
Will the test still pass? Could this be replaced with:
`match= "Job with name"`
This will assert the same behavior without a risk of a trivial change
breaking the test.
##########
providers/google/tests/unit/google/cloud/operators/test_gen_ai.py:
##########
@@ -435,6 +598,65 @@ def test_execute(self, mock_hook):
job_name=TEST_BATCH_JOB_NAME,
)
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def test_execute_value_error_raises_airflow_exception(self, mock_hook):
+ op = GenAIGeminiDeleteBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ job_name=TEST_BATCH_JOB_NAME,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ mock_hook.return_value.delete_batch_job.side_effect = ValueError()
+
+ with pytest.raises(
+ AirflowException, match=f"('Job with name %s was not found',
'{TEST_BATCH_JOB_NAME}')"
Review Comment:
The above comment applies to this one as well.
##########
providers/google/tests/unit/google/cloud/operators/test_gen_ai.py:
##########
@@ -458,6 +680,35 @@ def test_execute(self, mock_hook):
job_name=TEST_BATCH_JOB_NAME,
)
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def test_execute_propagates_value_error(self, mock_hook):
+ op = GenAIGeminiCancelBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ job_name=TEST_BATCH_JOB_NAME,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ mock_hook.return_value.cancel_batch_job.side_effect = ValueError()
+
+ with pytest.raises(
+ AirflowException, match=f"('Job with name %s was not found',
'{TEST_BATCH_JOB_NAME}')"
Review Comment:
Here as well.
I would suggest that you go through all the assertions of this type in the
PR and check for the same issue.
##########
providers/google/tests/unit/google/cloud/operators/test_gen_ai.py:
##########
@@ -435,6 +598,65 @@ def test_execute(self, mock_hook):
job_name=TEST_BATCH_JOB_NAME,
)
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def test_execute_value_error_raises_airflow_exception(self, mock_hook):
+ op = GenAIGeminiDeleteBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ job_name=TEST_BATCH_JOB_NAME,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ mock_hook.return_value.delete_batch_job.side_effect = ValueError()
+
+ with pytest.raises(
+ AirflowException, match=f"('Job with name %s was not found',
'{TEST_BATCH_JOB_NAME}')"
+ ):
+ op.execute(context={"ti": mock.MagicMock()})
+
+ mock_hook.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ )
+
+ mock_hook.return_value.delete_batch_job.assert_called_once_with(
+ job_name=TEST_BATCH_JOB_NAME,
+ )
+
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def test_execute_job_error_raises_airflow_exception(self, mock_hook):
+ op = GenAIGeminiDeleteBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ job_name=TEST_BATCH_JOB_NAME,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ mock_hook.return_value.delete_batch_job.return_value =
mock.MagicMock(error="Test error")
+
+ with pytest.raises(
+ AirflowException,
+ match=f"('Job with name %s was not deleted due to error: %s',
'{TEST_BATCH_JOB_NAME}', 'Test error')",
Review Comment:
Same here.
##########
providers/google/tests/unit/google/cloud/operators/test_gen_ai.py:
##########
@@ -152,6 +165,36 @@ def test_execute(self, mock_hook):
config=None,
)
+ @mock.patch(GEN_AI_PATH.format("GenAIGenerativeModelHook"))
+ def test_execute_propagates_client_error(self, mock_hook,
mock_client_error):
+ op = GenAIGenerateEmbeddingsOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ contents=CONTENTS,
+ model=EMBEDDING_MODEL,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ mock_hook.return_value.embed_content.side_effect = mock_client_error
+
+ with pytest.raises(ClientError):
+ op.execute(context={"ti": mock.MagicMock()})
+
+ mock_hook.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ mock_hook.return_value.embed_content.assert_called_once_with(
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ contents=CONTENTS,
+ model=EMBEDDING_MODEL,
+ config=None,
+ )
+
class TestGenAIGenerateContentOperator:
Review Comment:
Looks like there might be overlap here with PR #61262. The tests in that PR
look slightly different but I would just leave this out if they are testing the
same operator. But I would check with @shahar1 first.
--
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]
