Ghostkeeper Ghostkeeper - 3 months ago 14
Python Question

How to patch method io.RawIOBase.read with unittest?

I've recently learned about

unittest.monkey.patch
and its variants, and I'd like to use it to unit test for atomicity of a file read function. However, the patch doesn't seem to have any effect.

Here's my set-up. The method under scrutiny is roughly like so (abriged):

#local_storage.py

def read(uri):
with open(path, "rb") as file_handle:
result = file_handle.read()
return result


And the module that performs the unit tests (also abriged):

#test/test_local_storage.py

import unittest.mock
import local_storage

def _read_while_writing(io_handle, size=-1):
""" The patch function, to replace io.RawIOBase.read. """

_write_something_to(TestLocalStorage._unsafe_target_file) #Appends "12".
result = io_handle.read(size) #Should call the actual read.
_write_something_to(TestLocalStorage._unsafe_target_file) #Appends "34".

class TestLocalStorage(unittest.TestCase):
_unsafe_target_file = "test.txt"

def test_read_atomicity(self):
with open(self._unsafe_target_file, "wb") as unsafe_file_handle:
unsafe_file_handle.write(b"Test")

with unittest.mock.patch("io.RawIOBase.read", _read_while_writing): # <--- This doesn't work!
result = local_storage.read(TestLocalStorage._unsafe_target_file) #The actual test.
self.assertIn(result, [b"Test", b"Test1234"], "Read is not atomic.")


This way, the patch should ensure that every time you try to read it, the file gets modified just before and just after the actual read, as if it happens concurrently, thus testing for atomicity of our read.

The unit test currently succeeds, but I've verified with print statements that the patch function doesn't actually get called, so the file never gets the additional writes (it just says "Test"). I've also modified the code as to be non-atomic on purpose.

So my question: How can I patch the
read
function of an IO handle inside the local_storage module?
I've read elsewhere that people tend to replace the open() function to return something like a
StringIO
, but I don't see how that could fix this problem.

I need to support Python 3.4 and up.

Answer

I've finally found a solution myself.

The problem is that mock can't mock any methods of objects that are written in C. One of these is the RawIOBase that I was encountering.

So indeed the solution was to mock open to return a wrapper around RawIOBase. I couldn't get mock to produce a wrapper for me, so I implemented it myself.

There is one pre-defined file that's considered "unsafe". The wrapper writes to this "unsafe" file every time any call is made to the wrapper. This allows for testing the atomicity of file writes, since it writes additional things to the unsafe file while writing. My implementation prevents this by writing to a temporary ("safe") file and then moving that file over the target file.

The wrapper has a special case for the read function, because to test atomicity properly it needs to write to the file during the read. So it reads first halfway through the file, then stops and writes something, and then reads on. This solution is now semi-hardcoded (in how far is halfway), but I'll find a way to improve that.

You can see my solution here: https://github.com/Ghostkeeper/Luna/blob/0e88841d19737fb1f4606917f86e3de9b5b9f29b/plugins/storage/localstorage/test/test_local_storage.py