diff options
| -rw-r--r-- | tests/mock/disk.py | 17 |
1 files changed, 16 insertions, 1 deletions
diff --git a/tests/mock/disk.py b/tests/mock/disk.py index 39af5e8c3..1ff2a5c50 100644 --- a/tests/mock/disk.py +++ b/tests/mock/disk.py @@ -36,6 +36,18 @@ class DiskIO(object): def close(self): self.flush() return StringIO.close(self) + + def __del__(self): + try: + self.close() + except (AttributeError): + pass + + def __enter__(self): + return self + + def __exit__(self, *_): + self.close() class Dir(object): pass @@ -65,6 +77,9 @@ class DiskIO(object): content = self.fs.get(path, None) if content == self.Dir: raise IOError("[Errno 21] Is a directory: '%s'" % (path)) + elif mode == "w": + self.fs[path] = "" + f = self.TestFile(self.fs, path, self.fs[path]) elif not content and mode.startswith("w"): self.fs[path] = "" f = self.TestFile(self.fs, path, self.fs[path]) @@ -74,7 +89,7 @@ class DiskIO(object): f = self.TestFile(self.fs, path, content) f.seek(0, os.SEEK_END) else: - f = self.IO(self.fs, path, content) + f = self.TestFile(self.fs, path, content) return f |