如何为调用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

