用 Python 测试 Parquet

蔡斯哈丁

我试图模拟parquet并断言它是用正确的路径调用的，但无法正确模拟它。我怎样才能模拟这个option函数来返回一个模拟的parquet？

被测代码

def read_from_s3(spark, path):
    return spark.read.option('mergeSchema', 'true').parquet(path)

测试

import unittest
import mock
from src.read_from_s3 import read_from_s3


class TestReadFromS3(unittest.TestCase):
    def test_read_from_s3__called_with_correct_params(self):
        spark = mock.MagicMock()
        spark.read.option = mock.MagicMock()
        spark.read.option.parquet = mock.MagicMock()
        path = 'my_path'

        read_from_s3(spark, path)

        spark.read.option.assert_called_with('mergeSchema', 'true') # this passes
        spark.read.option.parquet.assert_called_with(path) # this fails

测试失败

AssertionError: Expected 'parquet' to have been called once. Called 0 times.

切普纳

parquet不是的属性spark.read.option；它的的属性返回值的option。此外，您不需要显式创建模拟；模拟上的属性查找也返回模拟。

def test_read_from_s3__called_with_correct_params(self):
    spark = mock.MagicMock()
    read_from_s3(spark, path)
    spark.read.option.assert_called_with('mergeSchema', 'true')
    spark.read.option.return_value.parquet.assert_called_with(path)

本文收集自互联网，转载请注明来源。

如有侵权，请联系 [email protected] 删除。