unittest unable to import class from pickle (AttributeError: Can't get attribute...)

Question

我需要一个单元测试来加载之前保存的 class pickle。但是，当我在 unittest 中加载泡菜时（在 unittest 工作之外），它会引发错误：

AttributeError: Can't get attribute 'Foo' on <module 'unittest.main' from '...\unittest\main.py'>

保存class的代码示例（我将这段代码保存在run_and_save_class.py）：

from pickle import dump
from pickle import load
from pickle import HIGHEST_PROTOCOL

class Foo(object):
    
    def __init__(self):
        self.bar = None
        self.file_out = "./out.pkl"

    def save_class(self):
        with open(self.file_out, "wb") as file_out:
            dump(self, file_out, protocol=HIGHEST_PROTOCOL)
            
    def load_class(self):
        with open(self.file_out, "rb") as file_out:
            cls = load(file_out)
        return cls

if __name__ == "__main__":
    cls = Foo()
    cls.bar = "saving a bar"
    cls.save_class()

用于测试 class 的代码（我将此代码保存在 unittest_class.py 中）：

import unittest
from run_and_save_class import Foo


class ClassValidation(unittest.TestCase):
    
    def __init__(self, *args, **kwargs):
        print("init")
        self.cls = Foo
        self.instance = Foo().load_class()
        print("class loaded")
        unittest.TestCase.__init__(self, *args, **kwargs)
    
    def test_anything(self):
        pass

我在 Anaconda 提示符中运行：

python run_and_save_class.py

python -m unittest -v unittest_class.py

后者会引发错误。

但是，这在笔记本中有效。

from run_and_save_class import Foo
cls = Foo().load_class()

我不明白为什么它不在单元测试中。

Answer 1

尝试以这种方式在类验证中实例化 Foo 对象：self.cls = Foo()

Answer 2

问题是 pickle 保存了相对于 __main__ 的对象，其中 dump 被调用（通过 save_class）。要加载相同的对象，您必须提供相同的环境 - 解决方法是在测试中将 class 添加到 __main__，以便 pickle 可以找到它：

import __main__

class ClassValidation(unittest.TestCase):

    def __init__(self, *args, **kwargs):
        __main__.Foo = Foo
        self.cls = Foo
        self.instance = Foo().load_class()
        unittest.TestCase.__init__(self, *args, **kwargs)

    def test_anything(self):
        self.assertEqual("saving a bar", self.instance.bar)

unittest unable to import class from pickle (AttributeError: Can't get attribute...)

unittest unable to import class from pickle (AttributeError: Can't get attribute...)

python

unit-testing

class

python-unittest

python-3.6