1 files changed, 27 insertions, 3 deletions
diff --git a/test/units/mock/loader.py b/test/units/mock/loader.py
index 9dc32cae..f6ceb379 100644
--- a/test/units/mock/loader.py
+++ b/test/units/mock/loader.py
@@ -21,15 +21,16 @@ __metaclass__ = type
 
 import os
 
+from ansible.errors import AnsibleParserError
 from ansible.parsing.dataloader import DataLoader
-from ansible.module_utils.common.text.converters import to_bytes, to_text
+from ansible.module_utils._text import to_bytes, to_text
 
 
 class DictDataLoader(DataLoader):
 
     def __init__(self, file_mapping=None):
         file_mapping = {} if file_mapping is None else file_mapping
-        assert isinstance(file_mapping, dict)
+        assert type(file_mapping) == dict
 
         super(DictDataLoader, self).__init__()
 
@@ -47,7 +48,11 @@ class DictDataLoader(DataLoader):
     # TODO: the real _get_file_contents returns a bytestring, so we actually convert the
     #       unicode/text it's created with to utf-8
     def _get_file_contents(self, file_name):
-        return to_bytes(self._file_mapping[file_name]), False
+        path = to_text(file_name)
+        if path in self._file_mapping:
+            return to_bytes(self._file_mapping[file_name]), False
+        else:
+            raise AnsibleParserError("file not found: %s" % file_name)
 
     def path_exists(self, path):
         path = to_text(path)
@@ -86,6 +91,25 @@ class DictDataLoader(DataLoader):
                 self._add_known_directory(dirname)
                 dirname = os.path.dirname(dirname)
 
+    def push(self, path, content):
+        rebuild_dirs = False
+        if path not in self._file_mapping:
+            rebuild_dirs = True
+
+        self._file_mapping[path] = content
+
+        if rebuild_dirs:
+            self._build_known_directories()
+
+    def pop(self, path):
+        if path in self._file_mapping:
+            del self._file_mapping[path]
+            self._build_known_directories()
+
+    def clear(self):
+        self._file_mapping = dict()
+        self._known_directories = []
+
     def get_basedir(self):
         return os.getcwd()