You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

如何为调用Google Cloud Logging API的fetch_remote_logs函数编写单元测试(使用Mock/Patch)

如何为调用Google Cloud Logging API的fetch_remote_logs函数编写单元测试(使用Mock/Patch)

你好呀!我不是软件工程师,这些测试相关的概念对我来说都是新的。我正在给一个仓库提交PR来优化功能,功能本身已经正常工作了,但现在需要为它写单元测试——这正是我不熟悉的部分。这个测试要针对的函数需要调用Google Cloud的库,我不太明白该怎么用Mock或者Patch来替换Cloud Logging API的请求结果。

为了让你更好理解我的问题,下面是我写的fetch_remote_logs函数:

class DbtGcpCloudRunJobBaseOperator(AbstractDbtBase, CloudRunExecuteJobOperator):  # type: ignore
    """
    Executes a dbt core cli command in a Cloud Run Job instance with dbt installed in it.

    """

    template_fields: Sequence[str] = tuple(
        list(AbstractDbtBase.template_fields) + list(CloudRunExecuteJobOperator.template_fields)
    )

    intercept_flag = False

    def __init__(
        self,
        # arguments required by CloudRunExecuteJobOperator
        project_id: str,
        region: str,
        job_name: str,
        #
        profile_config: ProfileConfig | None = None,
        command: list[str] | None = None,
        environment_variables: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        self.profile_config = profile_config
        self.command = command
        self.environment_variables = environment_variables or DEFAULT_ENVIRONMENT_VARIABLES
        super().__init__(project_id=project_id, region=region, job_name=job_name, **kwargs)
        # In PR #1474, we refactored cosmos.operators.base.AbstractDbtBase to remove its inheritance from BaseOperator
        # and eliminated the super().__init__() call. This change was made to resolve conflicts in parent class
        # initializations while adding support for ExecutionMode.AIRFLOW_ASYNC. Operators under this mode inherit
        # Airflow provider operators that enable deferrable SQL query execution. Since super().__init__() was removed
        # from AbstractDbtBase and different parent classes require distinct initialization arguments, we explicitly
        # initialize them (including the BaseOperator) here by segregating the required arguments for each parent class.
        kwargs.update(
            {
                "project_id": project_id,
                "region": region,
                "job_name": job_name,
                "command": command,
                "environment_variables": environment_variables,
            }
        )
        base_operator_args = set(inspect.signature(CloudRunExecuteJobOperator.__init__).parameters.keys())
        base_kwargs = {}
        for arg_key, arg_value in kwargs.items():
            if arg_key in base_operator_args:
                base_kwargs[arg_key] = arg_value
        base_kwargs["task_id"] = kwargs["task_id"]
        CloudRunExecuteJobOperator.__init__(self, **base_kwargs)


    def fetch_remote_logs(self, result: Any) -> Any:
        """
        Fetch Google Cloud Run job logs from Google Cloud Logging to Airflow logs
        The function returns the a list of the log messages

        result: result from build and run command in Cloud Run Job
        """

        self.log.info("Attempt to retrieve logs from Google Cloud Logging")

        execution_name = result["latest_created_execution"]["name"]
        execution_time = result["latest_created_execution"]["create_time"]
        filter_ = f'resource.type = "cloud_run_job" AND resource.labels.job_name = "{self.job_name}" AND timestamp>="{execution_time}"'

        time.sleep(5)  # Add sleep time to make sure all the job logs are available when we do the request

        try:
            client = logging.Client(project=self.project_id)
            # Fetch logs associated with the job_name
            entries = client.list_entries(filter_=filter_)
            self.log.info(f"Listing logs of the execution {execution_name} at {execution_time}:")

            # List to store log messages
            log_messages = []

            for entry in entries:
                # Search for logs associated with the job executed
                if entry.labels["run.googleapis.com/execution_name"] == execution_name:
                    log_messages.append(entry.payload)
                    self.log.info(f"Cloud Run Log: {entry.payload}")

            return log_messages

        except GoogleCloudError as e:
            # Catch Google Cloud-related errors (e.g., permission issues)
            self.log.warning(f"Warning: Error retrieving logs from Google Cloud Logging: {str(e)}")
            return
            # Continue without raising an error, just log the warning

我目前尝试写的测试是这样的:

@skip_on_empty_operator
@patch("google.cloud.logging.Client.list_entries")
def test_dbt_gcp_cloud_run_job_fetch_clean_logs():
"""
Check that fetch remote logs run correctly.
"""
run_operator = DbtRunGcpCloudRunJobOperator(**BASE_KWARGS)
mock_list_entries = {
    "latest_created_execution": {"name": "test-job-cloud-run", "create_time": "2025-03-19T09:41:22.608063Z"}
}
log_messages = run_operator.fetch_remote_logs(mock_list_entries)

assert len(log_messages) > 0

我知道有时候我用词可能不太准确,如果有表达不对的地方欢迎纠正我,提前感谢你的任何建议!


别担心,咱们一步步来搞定这个测试!你的思路是对的——用patch来模拟Google Cloud Logging的客户端调用,但还有几个细节需要调整,我来帮你梳理清楚:

1. 正确传递Mock对象到测试函数

你用了@patch("google.cloud.logging.Client.list_entries"),但测试函数没有接收这个mock参数,需要把它加到函数参数里,这样才能控制它的返回值:

@skip_on_empty_operator
@patch("google.cloud.logging.Client.list_entries")
def test_dbt_gcp_cloud_run_job_fetch_clean_logs(mock_list_entries):
    # 后续代码...

2. 构造符合函数预期的Mock返回值

你的fetch_remote_logs函数会遍历client.list_entries()返回的结果,并且检查每个entry的labels["run.googleapis.com/execution_name"]是否匹配。所以需要模拟出包含对应属性的entry对象:

# 先创建模拟的entry对象
class MockLogEntry:
    def __init__(self, execution_name, payload):
        self.labels = {"run.googleapis.com/execution_name": execution_name}
        self.payload = payload

# 让mock返回我们构造的entry列表
mock_list_entries.return_value = [
    MockLogEntry("test-job-cloud-run", "dbt run started"),
    MockLogEntry("test-job-cloud-run", "dbt run completed successfully"),
    # 可以加一个不匹配的entry,测试过滤逻辑
    MockLogEntry("other-execution", "unrelated log")
]

3. 完善测试逻辑,验证过滤和返回结果

现在可以调用函数并验证结果是否符合预期:

@skip_on_empty_operator
@patch("google.cloud.logging.Client.list_entries")
def test_dbt_gcp_cloud_run_job_fetch_clean_logs(mock_list_entries):
    """Check that fetch remote logs run correctly."""
    # 初始化operator
    run_operator = DbtRunGcpCloudRunJobOperator(**BASE_KWARGS)
    # 构造传给fetch_remote_logs的result参数
    mock_result = {
        "latest_created_execution": {
            "name": "test-job-cloud-run",
            "create_time": "2025-03-19T09:41:22.608063Z"
        }
    }

    # 构造模拟的日志条目
    class MockLogEntry:
        def __init__(self, execution_name, payload):
            self.labels = {"run.googleapis.com/execution_name": execution_name}
            self.payload = payload

    # 设置mock的返回值
    mock_list_entries.return_value = [
        MockLogEntry("test-job-cloud-run", "dbt run started"),
        MockLogEntry("test-job-cloud-run", "dbt run completed successfully"),
        MockLogEntry("other-execution", "unrelated log")
    ]

    # 调用函数
    log_messages = run_operator.fetch_remote_logs(mock_result)

    # 验证结果:应该只包含匹配execution_name的两条日志
    assert len(log_messages) == 2
    assert "dbt run started" in log_messages
    assert "dbt run completed successfully" in log_messages
    # 验证mock被调用了,且使用了正确的filter参数
    expected_filter = (
        f'resource.type = "cloud_run_job" AND resource.labels.job_name = "{run_operator.job_name}" '
        f'AND timestamp>="2025-03-19T09:41:22.608063Z"'
    )
    mock_list_entries.assert_called_once_with(filter_=expected_filter)

4. 额外建议:跳过sleep提升测试速度

函数里的time.sleep(5)在测试时会拖慢速度,你可以patchtime.sleep来跳过等待:

from unittest.mock import patch

@skip_on_empty_operator
@patch("time.sleep")
@patch("google.cloud.logging.Client.list_entries")
def test_dbt_gcp_cloud_run_job_fetch_clean_logs(mock_list_entries, mock_sleep):
    # 测试代码...
    # 可选:验证sleep被调用了一次
    mock_sleep.assert_called_once_with(5)

这样整个测试就可以在不依赖真实Google Cloud服务的情况下运行,而且能验证你的函数逻辑是否正确:比如是否正确过滤了日志条目、是否正确构造了API的filter参数、是否处理了返回结果等。

备注:内容来源于stack exchange,提问作者aegn

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.04.13 19:19:34