如何用pytest测试KFP v2组件？Input/Output与Fixture问题求助

在使用pytest对KFP v2.ds1的Kubeflow组件做本地测试时（组件在流水线中可正常运行），处理Input/Output参数与pytest fixture时出现报错：调用组件后返回TaskOutputArgument对象而非预期的Dataset，无法访问path属性。

代码示例 #

1. 模拟数据集的pytest fixture（同时作为KFP组件） #

# ./fixtures/
@pytest.fixture
@component()
def sample_df(dataset: Output[Dataset]):
     df = pd.DataFrame(
         {
             'name': ['Ana', 'Maria', 'Josh'],
             'age': [15, 19, 22],
         }
     )
     dataset.path += '.csv'
     df.to_csv(dataset.path, index=False)
     return

2. 待测试的KFP组件（年龄翻倍功能） #

# ./src/
@component()
def double_ages(df_input: Input[Dataset], df_output: Output[Dataset]):
     df = pd.read_csv(df_input.path)
     
     double_df = df.copy()
     double_df['age'] = double_df['age']*2

     df_output.path += '.csv'
     double_df.to_csv(df_output.path, index=False)

3. 测试代码 #

#./tests/
@pytest.mark.usefixtures("sample_df")
def test_double_ages(sample_df):

    expected_df = pd.DataFrame(
        {
            'name': ['Ana', 'Maria', 'Josh'],
            'age': [30, 38, 44],
        }
    )

    df_component = double_ages(sample_df)
    df_output = df_component.outputs['df_output']
    df = pd.read_csv(df_output.path)
    
    assert df['age'].tolist() == expected_df['age'].tolist()

报错信息 #

AttributeError: 'TaskOutputArgument' object has no attribute 'path'

方法一：直接测试组件核心业务逻辑（推荐） #

KFP的@component装饰器会将函数包装为任务构造器，直接调用返回的是任务参数而非执行结果。最高效的方式是剥离核心逻辑单独测试：

步骤1：分离组件业务逻辑 #

修改待测试组件，把核心逻辑抽成独立函数：

# ./src/
def _double_ages_logic(df: pd.DataFrame) -> pd.DataFrame:
    double_df = df.copy()
    double_df['age'] = double_df['age'] * 2
    return double_df

@component()
def double_ages(df_input: Input[Dataset], df_output: Output[Dataset]):
     df = pd.read_csv(df_input.path)
     double_df = _double_ages_logic(df)
     df_output.path += '.csv'
     double_df.to_csv(df_output.path, index=False)

步骤2：编写逻辑测试 #

直接测试核心函数，无需依赖KFP的Input/Output机制：

# ./tests/
def test_double_ages_logic():
    sample_df = pd.DataFrame(
        {
            'name': ['Ana', 'Maria', 'Josh'],
            'age': [15, 19, 22],
        }
    )
    expected_df = pd.DataFrame(
        {
            'name': ['Ana', 'Maria', 'Josh'],
            'age': [30, 38, 44],
        }
    )
    
    result_df = _double_ages_logic(sample_df)
    assert result_df['age'].tolist() == expected_df['age'].tolist()

如果不想抽离逻辑，也可以通过组件函数.func访问原始函数，传入模拟的Input/Output对象：

# ./tests/
from kfp.dsl import Dataset, Input, Output
import tempfile
import os

def test_double_ages_direct():
    # 创建临时输入文件
    with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as f:
        sample_df = pd.DataFrame(
            {
                'name': ['Ana', 'Maria', 'Josh'],
                'age': [15, 19, 22],
            }
        )
        sample_df.to_csv(f.name, index=False)
    
    # 模拟Input/Output对象
    input_dataset = Dataset(path=f.name)
    output_dataset = Dataset(path=tempfile.mktemp(suffix='.csv'))
    
    # 直接调用原始组件函数
    double_ages.func(Input(input_dataset), Output(output_dataset))
    
    # 验证结果
    expected_df = pd.DataFrame(
        {
            'name': ['Ana', 'Maria', 'Josh'],
            'age': [30, 38, 44],
        }
    )
    result_df = pd.read_csv(output_dataset.path)
    assert result_df['age'].tolist() == expected_df['age'].tolist()
    
    # 清理临时文件
    os.unlink(f.name)
    os.unlink(output_dataset.path)

方法二：使用KFP本地运行器执行完整组件 #

如果需要测试完整的组件执行流程（包括KFP参数处理），可以用LocalPipelineRunner运行组件：

# ./tests/
from kfp.dsl import Pipeline
from kfp.local import LocalPipelineRunner

def test_double_ages_full():
    # 定义测试流水线
    @Pipeline
    def test_pipeline():
        sample_task = sample_df()
        double_task = double_ages(df_input=sample_task.outputs['dataset'])
        return double_task
    
    # 运行流水线并获取结果
    runner = LocalPipelineRunner()
    run_result = runner.run(test_pipeline())
    
    # 验证输出
    output_path = run_result.outputs['df_output'].path
    result_df = pd.read_csv(output_path)
    
    expected_df = pd.DataFrame(
        {
            'name': ['Ana', 'Maria', 'Josh'],
            'age': [30, 38, 44],
        }
    )
    assert result_df['age'].tolist() == expected_df['age'].tolist()

注意：此方法需避免pytest fixture与KFP组件装饰器的冲突，若fixture调用异常，可在测试函数内直接定义sample_df组件。

内容的提问来源于stack exchange，提问作者Gabriel Caldas