From 06f372e56d4552bbf80367bcac09e73e2a304529 Mon Sep 17 00:00:00 2001
From: modulos <modulos@protonmail.com>
Date: Thu, 1 Jun 2017 22:25:10 +0200
Subject: [PATCH] Fix tests, add remove_public_key

The tests for VitualMachineSerializer were wrong. The manager now
contains a function to remove public keys from a user
---
 opennebula_api/models.py | 42 +++++++++++++++++++++++++++++++++
 opennebula_api/tests.py  | 51 ++++++++++------------------------------
 2 files changed, 54 insertions(+), 39 deletions(-)

diff --git a/opennebula_api/models.py b/opennebula_api/models.py
index 7872f6fc..4fa5d011 100644
--- a/opennebula_api/models.py
+++ b/opennebula_api/models.py
@@ -479,3 +479,45 @@ class OpenNebulaManager():
 
         except ConnectionError:
             raise
+
+    def remove_public_key(self, user, public_key=''):
+        """ 
+
+        Args: 
+            user (CustomUser): Dynamicweb user 
+            public_key (string): Public key to be removed to the user
+
+        Raises:
+            KeyDoesNotExistsError: If replace is False and the user already has a
+                public key 
+            WrongNameError: If no openebula user with this credentials exists
+            ConnectionError: If the connection to the opennebula server can't be
+                established 
+
+        Returns:
+            True if public_key was removed
+
+        """
+
+        try:
+            open_user = self._get_user(user)
+            try:
+                old_key = open_user.template.ssh_public_key 
+                if public_key not in old_key:
+                    raise KeyDoesNotExistsError()
+                if '\n{}'.format(public_key) in old_key:
+                    public_key = old_key.replace('\n{}'.format(public_key), '')
+                else: 
+                    public_key = old_key.replace(public_key, '')
+
+            except AttributeError:
+                raise KeyDoesNotExistsError()
+                
+            self.oneadmin_client.call('user.update', open_user.id,
+                         '<CONTEXT><SSH_PUBLIC_KEY>{key}</SSH_PUBLIC_KEY></CONTEXT>'.format(key=public_key))
+            return True
+        except WrongNameError:
+            raise
+
+        except ConnectionError:
+            raise
diff --git a/opennebula_api/tests.py b/opennebula_api/tests.py
index 073eb0b9..12b396fe 100644
--- a/opennebula_api/tests.py
+++ b/opennebula_api/tests.py
@@ -5,6 +5,7 @@ import string
 from django.test import TestCase
 
 from .models import OpenNebulaManager
+from .serializers import VirtualMachineSerializer
 from utils.models import CustomUser
 
 class OpenNebulaManagerTestCases(TestCase):
@@ -105,6 +106,8 @@ class OpenNebulaManagerTestCases(TestCase):
         user = self.manager._get_user(self.user)
         public_key = 'test'
         self.manager.add_public_key(self.user, public_key)
+        self.manager.add_public_key(self.user, public_key, merge=True)
+        user.info()
         old_public_key = user.template.ssh_public_key
         self.manager.remove_public_key(self.user, public_key)
         user.info()
@@ -112,7 +115,8 @@ class OpenNebulaManagerTestCases(TestCase):
         # Cleanup 
         user.delete()
 
-        self.assertEqual(new_public_key, old_public_key.replace(public_key, ''))
+        self.assertEqual(new_public_key,
+                old_public_key.replace('{}\n'.format(public_key), '', 1))
 
 
     def test_requires_ssh_key_for_new_vm(self):
@@ -120,49 +124,18 @@ class OpenNebulaManagerTestCases(TestCase):
         creating a new vm"""
 
 
-class VirtualMachineTestCase(TestCase):
+class VirtualMachineSerializerTestCase(TestCase):
     def setUp(self):
         """Define the test client and other test variables."""
-        self.template_name = "Standard"
-        self.base_price = 0.0
-        self.core_price = 5.0
-        self.memory_price = 2.0
-        self.disk_size_price = 0.6
-
-        self.cores = 1 
-        self.memory = 1
-        self.disk_size = 10.0
-        self.manager = OpenNebulaManager(email=None, password=None, create_user=False)
-        self.opennebula_id = self.manager.create_template(name=self.template_name,
-                                                          cores=self.cores,
-                                                          memory=self.memory,
-                                                          disk_size=self.disk_size)
-
-        self.template = VirtualMachineTemplate(opennebula_id=self.opennebula_id,
-                                               base_price=self.base_price,
-                                               memory_price=self.memory_price,
-                                               core_price=self.core_price,
-                                               disk_size_price=self.disk_size_price)
-        self.template_id = self.template.opennebula_id()
-        self.opennebula_id = self.manager.create_virtualmachine(template_id=self.template_id)
+        self.manager = OpenNebulaManager(email=None, password=None)
                                            
-        self.virtualmachine = VirtualMachine(opennebula_id=self.opennebula_id,
-                                             template=self.template)
 
     def test_serializer_strips_of_public(self):
-        """ Test the serialized object contains no 'public-'.""" 
+        """ Test the serialized virtual machine object contains no 'public-'.""" 
 
-        template = self.manager.get_templates().first()
-        serialized = VirtualMachineTemplateSerializer(template)
-        self.assertEqual(serialized.data.name, template.name.strip('public-'))
-
-
-        
-    def test_model_can_create_a_virtualmachine(self):
-        """Test the virtualmachine model can create a virtualmachine."""
-        old_count = VirtualMachine.objects.count()
-        self.virtualmachine.save()
-        new_count = VirtualMachine.objects.count()
-        self.assertNotEqual(old_count, new_count)
+        for vm in self.manager.get_vms():
+            serialized = VirtualMachineSerializer(vm)
+            self.assertEqual(serialized.data.get('name'), vm.name.strip('public-'))
+            break