1import unittest
2import textwrap
3import copy
4import pickle
5import email
6import email.message
7from email import policy
8from email.headerregistry import HeaderRegistry
9from test.test_email import TestEmailBase, parameterize
10
11
12@parameterize
13class TestPickleCopyHeader(TestEmailBase):
14
15    header_factory = HeaderRegistry()
16
17    unstructured = header_factory('subject', 'this is a test')
18
19    header_params = {
20        'subject': ('subject', 'this is a test'),
21        'from':    ('from',    'frodo@mordor.net'),
22        'to':      ('to',      'a: k@b.com, y@z.com;, j@f.com'),
23        'date':    ('date',    'Tue, 29 May 2012 09:24:26 +1000'),
24        }
25
26    def header_as_deepcopy(self, name, value):
27        header = self.header_factory(name, value)
28        h = copy.deepcopy(header)
29        self.assertEqual(str(h), str(header))
30
31    def header_as_pickle(self, name, value):
32        header = self.header_factory(name, value)
33        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
34            p = pickle.dumps(header, proto)
35            h = pickle.loads(p)
36            self.assertEqual(str(h), str(header))
37
38
39@parameterize
40class TestPickleCopyMessage(TestEmailBase):
41
42    # Message objects are a sequence, so we have to make them a one-tuple in
43    # msg_params so they get passed to the parameterized test method as a
44    # single argument instead of as a list of headers.
45    msg_params = {}
46
47    # Note: there will be no custom header objects in the parsed message.
48    msg_params['parsed'] = (email.message_from_string(textwrap.dedent("""\
49        Date: Tue, 29 May 2012 09:24:26 +1000
50        From: frodo@mordor.net
51        To: bilbo@underhill.org
52        Subject: help
53
54        I think I forgot the ring.
55        """), policy=policy.default),)
56
57    msg_params['created'] = (email.message.Message(policy=policy.default),)
58    msg_params['created'][0]['Date'] = 'Tue, 29 May 2012 09:24:26 +1000'
59    msg_params['created'][0]['From'] = 'frodo@mordor.net'
60    msg_params['created'][0]['To'] = 'bilbo@underhill.org'
61    msg_params['created'][0]['Subject'] = 'help'
62    msg_params['created'][0].set_payload('I think I forgot the ring.')
63
64    def msg_as_deepcopy(self, msg):
65        msg2 = copy.deepcopy(msg)
66        self.assertEqual(msg2.as_string(), msg.as_string())
67
68    def msg_as_pickle(self, msg):
69        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
70            p = pickle.dumps(msg, proto)
71            msg2 = pickle.loads(p)
72            self.assertEqual(msg2.as_string(), msg.as_string())
73
74
75if __name__ == '__main__':
76    unittest.main()
77