在python中对函数进行unittest

Question

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

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

函数中有两个变量，colinput和coloutput。

当我在新文件中运行unittest文件时，我正在导入函数(让我们将其命名为getweight()

)从mainscript中导入getweight

在getweight中，我计算了两个列计数

如何编写一个比较colinput == coloutput的测试?

我在如何从主脚本文件导入这些变量的部分卡住了。

Thanks in advance

Answer 1

将被测试的计算分解为一个纯函数。确保main是可导入的，没有副作用(使用'__main__'守卫):

main.py:

import pandas as pd
def getweight(x: pd.DataFrame) -> pd.DataFrame:
# ...
if __name__ == '__main__':
# run your computations

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

unittest.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)

* not testing