Reviewed-by: Bianca Carvalho <bianca@linux.vnet.ibm.com>

Signed-off-by: Ramon Medeiros <ramonn@linux.vnet.ibm.com>
---
 tests/test_model.py | 18 +++++++-----------
 tests/test_rest.py  | 10 ++++++++++
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/tests/test_model.py b/tests/test_model.py
index d9ffd5e..05f046c 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -325,25 +325,21 @@ class ModelTests(unittest.TestCase):
                   'source_media': {'type': 'disk', 'path': UBUNTU_ISO}}
         inst.templates_create(params)
         with RollbackContext() as rollback:
-            params = {'name': 'kimchi-vnc',
+            params = {'name': 'kimchi-graphics',
                       'template': '/plugins/kimchi/templates/test'}
             task1 = inst.vms_create(params)
             inst.task_wait(task1['id'])
-            rollback.prependDefer(inst.vm_delete, 'kimchi-vnc')
+            rollback.prependDefer(inst.vm_delete, 'kimchi-graphics')
 
-            info = inst.vm_lookup('kimchi-vnc')
+            info = inst.vm_lookup('kimchi-graphics')
             self.assertEquals('vnc', info['graphics']['type'])
             self.assertEquals('127.0.0.1', info['graphics']['listen'])
 
-            graphics = {'type': 'spice', 'listen': '127.0.0.1'}
-            params = {'name': 'kimchi-spice',
-                      'template': '/plugins/kimchi/templates/test',
-                      'graphics': graphics}
-            task2 = inst.vms_create(params)
-            inst.task_wait(task2['id'])
-            rollback.prependDefer(inst.vm_delete, 'kimchi-spice')
+            graphics = {'type': 'spice'}
+            params = {'graphics': graphics}
+            inst.vm_update('kimchi-graphics', params)
 
-            info = inst.vm_lookup('kimchi-spice')
+            info = inst.vm_lookup('kimchi-graphics')
             self.assertEquals('spice', info['graphics']['type'])
             self.assertEquals('127.0.0.1', info['graphics']['listen'])
 
diff --git a/tests/test_rest.py b/tests/test_rest.py
index 90a808a..76f0b9c 100644
--- a/tests/test_rest.py
+++ b/tests/test_rest.py
@@ -301,6 +301,16 @@ class RestTests(unittest.TestCase):
         resp = self.request('/plugins/kimchi/vms/∨м-црdαtеd', req, 'PUT')
         self.assertEquals(400, resp.status)
 
+        # change vm graphics type
+        req = json.dumps({"graphics": {"type":"spice"}})
+        resp = self.request('/plugins/kimchi/vms/∨м-црdαtеd', req, 'PUT')
+        self.assertEquals(json.loads(resp.read())["graphics"]["type"], "spice")
+
+        # try to add a invalid type
+        req = json.dumps({"graphics": {"type":"test"}})
+        resp = self.request('/plugins/kimchi/vms/∨м-црdαtеd', req, 'PUT')
+        self.assertEquals(400, resp.status)
+
     def test_vm_lifecycle(self):
         # Create a Template
         req = json.dumps({'name': 'test',