[Python][mock][HowTo] 替換內建的 open()

程式裡如果有對檔案做操作,免不了會直接用到 open() builtin,在做單元測試時要如何替換 (mock out) 掉呢?

mock 很貼心地提供了 mock_open() helper:

mock_open(mock=None, read_data='')

其中 mock 可以提供自訂的 mock object (否則內部自動生出一個 MagicMock),然後內部會將它 config 成像 open() 一樣,被呼叫時會傳回 file-like object,read_data 則可以進一步設定 file-like object 的 read() 要傳回什麼內容。用起來像是這樣子:

>>> from mock import mock_open
>>> open = mock_open(read_data='content') # 2
>>> f = open('filename', 'rb') # 1
>>> f.read()                              # 2
'content'
>>> f.read()
'content'
>>>
>>> with open('filename', 'rb') as f:     # 3
...     print f.read()
...
content
1 mock object 被呼叫時會傳回 file-like 的 mock object。
2 之後呼叫 file-like 的 read() 固定會傳回 read_data 參數傳入的內容。
3 mock_open() 也支援 context manager 的用法。

用實際的例子做說明;hashutil.py 裡的 md5sum_memory_inefficient()md5sum() 都可以用來計算檔案的 MD5 checksum,只是前者的寫法比較單純(但沒有效率),先來看怎麼利用 mock object 對它做單元測試:

hashutil.py
import hashlib

def md5sum_memory_inefficient(filename):
    with open(filename, 'rb') as f:
        return hashlib.md5(f.read()).hexdigest() # 1
1 f.read() 一次將整個檔案內容讀入的做法,遇到大檔案時很容易就會塞爆記憶體。
test_hashutil.py
import unittest2 as unittest
from mock import patch, mock_open
import hashutil

class HashTest(unittest.TestCase):

    @patch('hashutil.open', mock_open(read_data=b'content'), create=True) # 1
    def test_md5sum_memory_inefficient(self):
        self.assertEqual(hashutil.md5sum_memory_inefficient('filename'),
                         '9a0364b9e99bb480dd25e1f0284c8555')              # 1

        filelike = hashutil.open.return_value
        filelike.read.assert_called_once_with() # 2
1 read_data 提供的檔案內容 “content",其 MD5 digest 是 `9a0364b9e99bb480dd25e1f0284c8555
Caution 這裡替換 hashutil.open 而非 __builtin__.open,否則會影響到整個 Python runtime。按照 LEGB naming lookup 的原則,會由 global (module) 提供 mock open()

2由於是將檔案內容一次讀入,所以 read() 只會被呼叫一次。

接著來看另一個比較有效率的版本:(切割 chunks 分批讀入)

hashutil.py
import hashlib

def md5sum(filename):
    md5 = hashlib.md5()
    with open(filename, 'rb') as f:
        for chunk in iter(lambda: f.read(8192), ''):
            md5.update(chunk)
    return md5.hexdigest()

先用上述的方法測試:

    @patch('hashutil.open', mock_open(read_data=b'chunk1chunk2chunk3'), create=True) # 1
    def test_md5sum(self):
        self.assertEqual(hashutil.md5sum('filename'),
                         '2aca0a9378723b1bed59975523ed50cd')
1 這一次假設檔案很大,要切割成不同的 chunks 分批讀入。

但實際執行測試卻整個卡住了?那是因為 iter(lambda: f.read(8192), “) 這種較為 Pythonic 的寫法會不斷呼叫 read() 直到遇見 (空字串) 才會停下來。

因為 mock_open() 只是簡單設定了 read() 的回傳值,為了模擬這裡 chunks 分批讀入的狀況,必須要對 mock_open() 傳回的 mock object 做一點加工才行。

    @patch('hashutil.open', mock_open(), create=True)
    def test_md5sum(self):
        mock_open = hashutil.open # 1
        filelike = mock_open.return_value
        filelike.read.side_effect = [b'chunk1', b'chunk2', b'chunk3', b''] # 2

        self.assertEqual(hashutil.md5sum('filename'),
                         '2aca0a9378723b1bed59975523ed50cd')
        self.assertEqual(filelike.read.call_count, 4) # 3
1 取回 @patch 時用 mock_open() 產生的 mock object。
2 安排 file-like object 的 read() 依序傳回 chunks,最後的 b" 表示 EOF。
3 驗證 read() 確實被呼叫了 4 次,並非一次傳回整個檔案的內容。

總結一下

  • mock_open() 產生的 mock objects 可以直接替換 read() 一次將檔案內容讀入的做法。
  • 如果實作是採用切割 chunks 分批讀入的做法,則要另外安排 file-like object 其 read() 會連續被呼叫的行為。
廣告

發表迴響

Please log in using one of these methods to post your comment:

WordPress.com Logo

您的留言將使用 WordPress.com 帳號。 登出 / 變更 )

Twitter picture

您的留言將使用 Twitter 帳號。 登出 / 變更 )

Facebook照片

您的留言將使用 Facebook 帳號。 登出 / 變更 )

Google+ photo

您的留言將使用 Google+ 帳號。 登出 / 變更 )

連結到 %s