From 1c486a86812d671a773d54cba23fd0b21a7dc49e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Janne=20R=C3=B6nkk=C3=B6?= <janne.ronkko@vincit.fi>
Date: Thu, 2 Feb 2017 15:56:08 +0200
Subject: [PATCH] Fix disk usage calculations when too large file is created

Previously if you created file having larger initial contents than the
free space in the fake filesystem, the disk usage counter was decreased
by the size of the initial contents possibly resulting in negative usage.

Basically the issue was that when ChangeDiskUsage is called with size
larger than the amount of free space, space usage is not changed but
an error is raised. Due to an exception, CreateFile removes the created
object that results in new call to ChangeDiskUsage with negative size
of the object.

Previously the object's size was set before ChangeDiskUsage was called
and then when the object was being removed the object's already set size
was used to reduce the filesystem usage (even though that size was never
added to the usage).
---
 fake_filesystem_test.py     | 35 +++++++++++++++++++++++++++++++++++
 pyfakefs/fake_filesystem.py |  4 ++--
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/fake_filesystem_test.py b/fake_filesystem_test.py
index 0e0b17e0..ea6e0df1 100755
--- a/fake_filesystem_test.py
+++ b/fake_filesystem_test.py
@@ -4165,6 +4165,25 @@ def setUp(self):
         self.filesystem = fake_filesystem.FakeFilesystem(path_separator='!', total_size=100)
         self.os = fake_filesystem.FakeOsModule(self.filesystem)
 
+    def testDiskUsageOnFileCreation(self):
+        fake_open = fake_filesystem.FakeFileOpen(self.filesystem)
+
+        total_size = 100
+        self.filesystem.AddMountPoint('mount', total_size)
+        def create_too_large_file():
+            with fake_open('!mount!file', 'w') as dest:
+                dest.write('a' * (total_size + 1))
+
+        self.assertRaises((OSError, IOError), create_too_large_file)
+
+        self.assertEqual(0, self.filesystem.GetDiskUsage('!mount').used)
+
+        with fake_open('!mount!file', 'w') as dest:
+            dest.write('a' * total_size)
+
+        self.assertEqual(total_size, self.filesystem.GetDiskUsage('!mount').used)
+
+
     def testFileSystemSizeAfterLargeFileCreation(self):
         filesystem = fake_filesystem.FakeFilesystem(path_separator='!',
                                                     total_size=1024 * 1024 * 1024 * 100)
@@ -4203,29 +4222,45 @@ def testFileSystemSizeAfterDirectoryRemoval(self):
         self.assertEqual((100, 40, 60), self.filesystem.GetDiskUsage())
 
     def testCreatingFileWithFittingContent(self):
+        initial_usage = self.filesystem.GetDiskUsage()
+
         try:
             self.filesystem.CreateFile('!foo!bar', contents=b'a' * 100)
         except IOError:
             self.fail('File with contents fitting into disk space could not be written.')
 
+        self.assertEqual(initial_usage.used + 100, self.filesystem.GetDiskUsage().used)
+
     def testCreatingFileWithContentTooLarge(self):
         def create_large_file():
             self.filesystem.CreateFile('!foo!bar', contents=b'a' * 101)
 
+        initial_usage = self.filesystem.GetDiskUsage()
+
         self.assertRaises(IOError, create_large_file)
 
+        self.assertEqual(initial_usage, self.filesystem.GetDiskUsage())
+
     def testCreatingFileWithFittingSize(self):
+        initial_usage = self.filesystem.GetDiskUsage()
+
         try:
             self.filesystem.CreateFile('!foo!bar', st_size=100)
         except IOError:
             self.fail('File with size fitting into disk space could not be written.')
 
+        self.assertEqual(initial_usage.used + 100, self.filesystem.GetDiskUsage().used)
+
     def testCreatingFileWithSizeTooLarge(self):
+        initial_usage = self.filesystem.GetDiskUsage()
+
         def create_large_file():
             self.filesystem.CreateFile('!foo!bar', st_size=101)
 
         self.assertRaises(IOError, create_large_file)
 
+        self.assertEqual(initial_usage, self.filesystem.GetDiskUsage())
+
     def testResizeFileWithFittingSize(self):
         file_object = self.filesystem.CreateFile('!foo!bar', st_size=50)
         try:
diff --git a/pyfakefs/fake_filesystem.py b/pyfakefs/fake_filesystem.py
index 95bada08..dd62bdc8 100644
--- a/pyfakefs/fake_filesystem.py
+++ b/pyfakefs/fake_filesystem.py
@@ -308,10 +308,10 @@ def _SetInitialContents(self, contents, encoding):
         if self.byte_contents:
             self.SetSize(0)
         current_size = self.st_size or 0
+        if self.filesystem:
+            self.filesystem.ChangeDiskUsage(st_size - current_size, self.name, self.st_dev)
         self.byte_contents = contents
         self.st_size = st_size
-        if self.filesystem:
-            self.filesystem.ChangeDiskUsage(self.st_size - current_size, self.name, self.st_dev)
         self.epoch += 1
 
     def SetContents(self, contents, encoding=None):