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

Code before:

        mock_sync_client.assert_called_once()
        call_kwargs = mock_sync_client.call_args[1]
        assert call_kwargs["transport"] == "rest"

    @pytest.mark.asyncio
    @mock.patch("asyncio.to_thread")
    async def test_get_operation_rest_uses_to_thread(self, mock_to_thread):
        """Test that get_operation uses asyncio.to_thread for REST transport."""
        expected_operation = operations_pb2.Operation(name=OPERATION_NAME)
        mock_to_thread.return_value = expected_operation

        hook = CloudRunAsyncHook(transport="rest")
        mock_conn = mock.MagicMock()  # sync client
        hook.get_conn = mock.AsyncMock(return_value=mock_conn)

        result = await hook.get_operation(operation_name=OPERATION_NAME)

        mock_to_thread.assert_called_once_with(
            mock_conn.get_operation,
            operations_pb2.GetOperationRequest(name=OPERATION_NAME),
            timeout=120,
        )
        assert result == expected_operation

    def mock_get_operation(self, expected_operation):
        get_operation_mock = mock.AsyncMock()
        get_operation_mock.return_value = expected_operation
        return get_operation_mock

    def _dummy_get_credentials(self):
        pass


@pytest.mark.db_test
class TestCloudRunServiceHook:
    def dummy_get_credentials(self):
        pass

    @pytest.fixture
    def cloud_run_service_hook(self):
        cloud_run_service_hook = CloudRunServiceHook()
        cloud_run_service_hook.get_credentials = self.dummy_get_credentials
        return cloud_run_service_hook

    @mock.patch(
        "airflow.providers.google.common.hooks.base_google.GoogleBaseAsyncHook.__init__",
        new=mock_base_gcp_hook_default_project_id,
    )
    @mock.patch("airflow.providers.google.cloud.hooks.cloud_run.ServicesClient")
    def test_get_service(self, mock_batch_service_client, cloud_run_service_hook):

Review comment (suggestion):
type the expected_operation as operations_pb2.Operation and also add spec to get_operation AsyncMock(spec=operations_pb2.Operation)

Code after:
        await hook.get_conn()

        mock_sync_client.assert_called_once()
        call_kwargs = mock_sync_client.call_args[1]
        assert call_kwargs["transport"] == "rest"

    @pytest.mark.asyncio
    @mock.patch("asyncio.to_thread")
    async def test_get_operation_rest_uses_to_thread(self, mock_to_thread):
        """Test that get_operation uses asyncio.to_thread for REST transport."""
        expected_operation = operations_pb2.Operation(name=OPERATION_NAME)
        mock_to_thread.return_value = expected_operation

        hook = CloudRunAsyncHook(transport="rest")
        mock_conn = mock.MagicMock(spec=JobsClient)  # sync client
        hook.get_conn = mock.AsyncMock(return_value=mock_conn)

        result = await hook.get_operation(operation_name=OPERATION_NAME)

        mock_to_thread.assert_called_once_with(
            mock_conn.get_operation,
            operations_pb2.GetOperationRequest(name=OPERATION_NAME),
            timeout=120,
        )
        assert result == expected_operation


@pytest.mark.db_test
class TestCloudRunServiceHook:
    def dummy_get_credentials(self):
        pass

    @pytest.fixture
    def cloud_run_service_hook(self):
        cloud_run_service_hook = CloudRunServiceHook()
        cloud_run_service_hook.get_credentials = self.dummy_get_credentials
        return cloud_run_service_hook

    @mock.patch(
        "airflow.providers.google.common.hooks.base_google.GoogleBaseAsyncHook.__init__",
        new=mock_base_gcp_hook_default_project_id,
    )
    @mock.patch("airflow.providers.google.cloud.hooks.cloud_run.ServicesClient")
    def test_get_service(self, mock_batch_service_client, cloud_run_service_hook):
        get_service_request = GetServiceRequest(
            name=f"projects/{PROJECT_ID}/locations/{REGION}/services/{SERVICE_NAME}"
        )

        cloud_run_service_hook.get_service(service_name=SERVICE_NAME, region=REGION, project_id=PROJECT_ID)
        cloud_run_service_hook._client.get_service.assert_called_once_with(get_service_request)