Repository: apache/airflow
File: providers/google/tests/unit/google/cloud/operators/test_gen_ai.py (Python)

Code before:
        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.cancel_batch_job.assert_called_once_with(
            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}')"
        ):
            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.cancel_batch_job.assert_called_once_with(
            job_name=TEST_BATCH_JOB_NAME,
        )


class TestGenAIGeminiCreateEmbeddingsBatchJobOperator:
    @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
    def test_execute(self, mock_hook):
        op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
            task_id=TASK_ID,
            project_id=GCP_PROJECT,
            location=GCP_LOCATION,
            input_source=TEST_EMBEDDINGS_JOB_INLINED_REQUESTS,
            model=EMBEDDING_MODEL,
            gemini_api_key=TEST_GEMINI_API_KEY,
            gcp_conn_id=GCP_CONN_ID,

Review comment (suggestion):
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. 

Code after:
        )
        mock_hook.return_value.cancel_batch_job.assert_called_once_with(
            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 = 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):
            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.cancel_batch_job.assert_called_once_with(
            job_name=TEST_BATCH_JOB_NAME,
        )


class TestGenAIGeminiCreateEmbeddingsBatchJobOperator:
    @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
    def test_execute(self, mock_hook):
        op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
            task_id=TASK_ID,
            project_id=GCP_PROJECT,
            location=GCP_LOCATION,
            input_source=TEST_EMBEDDINGS_JOB_INLINED_REQUESTS,
            model=EMBEDDING_MODEL,
            gemini_api_key=TEST_GEMINI_API_KEY,
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,
            wait_until_complete=False,
        )
        op.execute(context={"ti": mock.MagicMock()})
        mock_hook.assert_called_once_with(
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,