From 7779f7f89552f06cb924339bc865fd6a9c84c840 Mon Sep 17 00:00:00 2001
From: Joshua Boniface <joshua@boniface.me>
Date: Mon, 11 Jun 2018 02:32:48 -0400
Subject: [PATCH] Try adding regular node state change tests

---
 NodeInstance.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/NodeInstance.py b/NodeInstance.py
index 047cf904..471777a1 100644
--- a/NodeInstance.py
+++ b/NodeInstance.py
@@ -41,6 +41,7 @@ class NodeInstance():
         self.domains_count = 0
         self.memused = 0
         self.memfree = 0
+        self.inflush = False
 
         # Zookeeper handlers for changed states
         @zk.DataWatch('/nodes/{}/state'.format(self.name))
@@ -95,6 +96,7 @@ class NodeInstance():
 
     # Flush all VMs on the host
     def flush(self):
+        self.inflush = True
         ansiiprint.echo('Flushing node "{}" of running VMs'.format(self.name), '', 'i')
         for dom_uuid in self.domain_list:
             most_memfree = 0
@@ -127,6 +129,9 @@ class NodeInstance():
             # Wait 1s between migrations
             time.sleep(1)
 
+        self.zk.set('/nodes/{}/state'.format(self.name), 'flushed'.encode('ascii'))
+        self.inflush = False
+
     def unflush(self):
         ansiiprint.echo('Restoring node {} to active service.'.format(self.name), '', 'i')
         self.zk.set('/nodes/{}/state'.format(self.name), 'start'.encode('ascii'))
@@ -182,6 +187,12 @@ class NodeInstance():
                         except:
                             pass
 
+        # toggle state management of this node
+        if self.state == 'flush':
+            self.flush()
+        if self.state == 'unflush':
+            self.unflush()
+
         # Set our information in zookeeper
         self.name = conn.getHostname()
         self.cpucount = conn.getCPUMap()[0]