在python中对一个函数进行单元测试

Question

我有一个从数据框中获取列数的函数。

运行函数后，我想确保列数与输入列数相同。

所以，我在函数中有两个变量，colinput 和 coloutput。

当我在新文件中运行 unittest 文件时，我正在导入该函数（我们称之为 getweight()

所以，从 mainscript 导入 getweight

在 getweight 里面我计算了两列计数

如何编写比较 colinput == coloutput 的测试？

我在如何从主脚本文件中导入这些变量方面陷入困境。

提前致谢

Answer 1

将被测计算分解为纯函数。 确保 main 可以导入而没有副作用（使用 '__main__' 防护）：

主要.py：

import pandas as pd

def getweight(x: pd.DataFrame) -> pd.DataFrame:
  # ...

if __name__ == '__main__':
  # run your computations

导入您感兴趣的主模块和测试属性：

单元测试.py：

import pandas as pd
import main
import pytest

@pytest.mark.parametrize(
  'x',
  [
    # specify interesting inputs ...
  ]
)
def test_getweight_preserves_column_count(x: pd.DataFrame):
  assert len(x.columns) == len(main.getweight(x).columns)

* 未经测试