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


class TestGenAIGeminiListBatchJobsOperator:
    @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
    def test_execute(self, mock_hook):
        op = GenAIGeminiListBatchJobsOperator(
            task_id=TASK_ID,
            project_id=GCP_PROJECT,
            location=GCP_LOCATION,
            gemini_api_key=TEST_GEMINI_API_KEY,
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,
        )

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

Code after:
            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,
        )

        result = op.execute(context={"ti": mock.MagicMock()})

        assert result == expected_return
        mock_job.model_dump.assert_called_once_with(mode="json")

    @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
    def test_execute_value_error_raises_airflow_exception(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):
            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.get_batch_job.assert_called_once_with(
            job_name=TEST_BATCH_JOB_NAME,
        )


class TestGenAIGeminiListBatchJobsOperator:
    @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
    def test_execute(self, mock_hook):
        op = GenAIGeminiListBatchJobsOperator(
            task_id=TASK_ID,
            project_id=GCP_PROJECT,
            location=GCP_LOCATION,
            gemini_api_key=TEST_GEMINI_API_KEY,
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,
        )