Unit Test: Mock the open() builtin function on a specific filename

unstko · unstko · commit 558b9c43e88d · 2016-12-13T22:40:31.000+01:00
- input.txt gets mocked by test_input
- All other files (like JSON) do not get mocked
diff --git a/lib/test.py b/lib/test.py
@@ -1,8 +1,64 @@
 import unittest
 from unittest.mock import mock_open, patch
+from contextlib import contextmanager
+from io import StringIO
 import solution
 
 
+class NotMocked(Exception):
+    def __init__(self, filename):
+        super(NotMocked, self).__init__(
+            "The file %s was opened, but not mocked." % filename)
+        self.filename = filename
+
+
+@contextmanager
+def mock_open(filename, contents=None, complain=True):
+    """Mock the open() builtin function on a specific filename
+.
+    Let execution pass through to open() on files different than
+    :filename:. Return a StringIO with :contents: if the file was
+    matched. If the :contents: parameter is not given or if it is None,
+    a StringIO instance simulating an empty file is returned.
+.
+    If :complain: is True (default), will raise an AssertionError if
+    :filename: was not opened in the enclosed block. A NotMocked
+    exception will be raised if open() was called with a file that was
+    not mocked by mock_open.
+
+    Source: https://mapleoin.github.io/perma/mocking-python-file-open
+    """
+    open_files = set()
+
+    def mock_file(*args):
+        if args[0] == filename:
+            f = StringIO(contents)
+            f.name = filename
+        else:
+            mocked_file.stop()
+            f = open(*args)
+            mocked_file.start()
+        open_files.add(f.name)
+        return f
+
+    mocked_file = patch('builtins.open', mock_file)
+    mocked_file.start()
+    try:
+        yield
+    except NotMocked as e:
+        if e.filename != filename:
+            raise
+    mocked_file.stop()
+    try:
+        open_files.remove(filename)
+    except KeyError:
+        if complain:
+            raise AssertionError("The file %s was not opened." % filename)
+    for f_name in open_files:
+        if complain:
+            raise NotMocked(f_name)
+
+
 class Test(unittest.TestCase):
     def __init__(self, *args, **kwargs):
         super(Test, self).__init__(*args, **kwargs)
@@ -17,8 +73,7 @@ def tearDown(self):
         print("")
 
     def execute_test(self, test_input):
-        mocked_open_function = mock_open(read_data=test_input)
-        with patch("builtins.open", mocked_open_function):
+        with mock_open('/input.txt', test_input, False):
             self.solution.calculate(True)
 
     def get_solution(self, nr):
