From 1e4b6c57b3358731fa52110eb68a3c55fc3f5e7f Mon Sep 17 00:00:00 2001
From: numirias <numirias@users.noreply.github.com>
Date: Tue, 19 Dec 2017 19:14:21 +0100
Subject: [PATCH] Fix restoring at root node

---
 plasma/node.py     | 16 +++++++++++-----
 tests/test_node.py | 46 +++++++++++++++++++++++++++++++---------------
 2 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/plasma/node.py b/plasma/node.py
index 7f1c60c..2552851 100644
--- a/plasma/node.py
+++ b/plasma/node.py
@@ -547,12 +547,13 @@ def restore(self, node):
         Try to add the node in a place where a node with the same payload
         has previously been.
         """
+        restorables = self.root.restorables
         try:
-            parent, idx, sizes, flip = self.root.restorables[node.payload]
+            parent, idx, sizes, fixed, flip = restorables[node.payload]
         except KeyError:
             raise NotRestorableError()
         if parent not in self.root:
-            # Don't restore at a parent that's not part of the tree anymore
+            # Don't try to restore if parent is not part of the tree anymore
             raise NotRestorableError()
         node.reset_size()
         if flip:
@@ -565,18 +566,23 @@ def restore(self, node):
             node.size = sizes[0]
             if len(sizes) == 2:
                 node.siblings[0].size = sizes[1]
-        del self.root.restorables[node.payload]
+        if not fixed:
+            node.reset_size()
+        del restorables[node.payload]
 
     def _save_restore_state(self):
         parent = self.parent
-        sizes = (self._size,)
+        sizes = (self.size,)
         flip = False
         if len(self.siblings) == 1:
+            # If there is only one node left in the container, we need to save
+            # its size too because the size will be lost.
             sizes += (self.siblings[0]._size,)  # pylint: disable=W0212
             if not self.parent.is_root:
                 flip = True
                 parent = self.siblings[0]
-        self.root.restorables[self.payload] = (parent, self.index, sizes, flip)
+        self.root.restorables[self.payload] = (parent, self.index, sizes,
+                                               self.fixed, flip)
 
     def move(self, direction):
         if self.is_root:
diff --git a/tests/test_node.py b/tests/test_node.py
index 9a4efee..5ccebfd 100644
--- a/tests/test_node.py
+++ b/tests/test_node.py
@@ -893,21 +893,6 @@ def test_resize_max(self, root, tiny_grid):
         assert a.width == 110
         assert b.width == c.width == 10
 
-    def test_resize_with_collapse_and_restore(self, root, small_grid):
-        a, b, c, d = small_grid
-        root.height = 30
-        c.size = 30
-        d.size += 10
-        b.remove()
-        assert c.size == c.height == 10
-        assert d.size == d.height == 20
-        root.restore(b)
-        assert b.height == 15
-        assert b.width == 60
-        assert c.height == d.height == 15
-        assert c.width == 20
-        assert d.width == 40
-
 class TestRestore:
 
     def test_restore(self, root, grid):
@@ -1001,6 +986,22 @@ def test_restore_root(self, root):
         assert a._size == 20
         assert b._size is None
 
+    def test_restore_root2(self, root):
+        a, b, c = Nodes('a b c')
+        root.add_child(a)
+        root.add_child(b)
+        root.add_child(c)
+        b.size = 20
+        c.size = 40
+        a.remove()
+        assert b.size == 40
+        assert c.size == 80
+        root.restore(a)
+        assert not a.fixed
+        assert a.size == 60
+        assert b.size == 20
+        assert c.size == 40
+
     def test_restore_keep_flexible(self, root, tiny_grid):
         a, b, c = tiny_grid
         b.remove()
@@ -1025,3 +1026,18 @@ def test_restore_keep_flexible(self, root, tiny_grid):
         root.restore(c)
         assert b._size is None
         assert c._size == 10
+
+    def test_resize_with_collapse_and_restore(self, root, small_grid):
+        a, b, c, d = small_grid
+        root.height = 30
+        c.size = 30
+        d.size += 10
+        b.remove()
+        assert c.size == c.height == 10
+        assert d.size == d.height == 20
+        root.restore(b)
+        assert b.height == 15
+        assert b.width == 60
+        assert c.height == d.height == 15
+        assert c.width == 20
+        assert d.width == 40