From 29f0be53fab51002248837d62efcd9381446ac57 Mon Sep 17 00:00:00 2001
From: Jitendra Yejare <jyejare@redhat.com>
Date: Mon, 5 Feb 2018 15:10:39 +0530
Subject: [PATCH] [6.3 Feature] SSH Key helpers added in User entity (#478)

---
 nailgun/entities.py    | 75 ++++++++++++++++++++++++++++++++++++++++++
 tests/test_entities.py | 21 ++++++++++++
 2 files changed, 96 insertions(+)

diff --git a/nailgun/entities.py b/nailgun/entities.py
index 5fe3582e..c74b7bbd 100644
--- a/nailgun/entities.py
+++ b/nailgun/entities.py
@@ -5941,6 +5941,81 @@ def update_payload(self, fields=None):
         }
 
 
+class SSHKey(
+        Entity,
+        EntityCreateMixin,
+        EntityDeleteMixin,
+        EntityReadMixin,
+        EntitySearchMixin):
+    """A representation of a SSH Key entity.
+
+    ``user`` must be passed in when this entity is instantiated.
+
+    :raises: ``TypeError`` if ``user`` is not passed in.
+
+    """
+
+    def __init__(self, server_config=None, **kwargs):
+        _check_for_value('user', kwargs)
+        self._fields = {
+            'user': entity_fields.OneToOneField(
+                User,
+                required=True,
+            ),
+            'name': entity_fields.StringField(
+                required=True,
+                str_type='alpha',
+                length=(6, 12),
+                unique=True
+            ),
+            'key': entity_fields.StringField(
+                required=True,
+                str_type='alphanumeric',
+                unique=True
+            )
+        }
+        super(SSHKey, self).__init__(server_config, **kwargs)
+        self._meta = {
+            # pylint:disable=no-member
+            'api_path': '{0}/ssh_keys'.format(self.user.path()),
+        }
+
+    def read(self, entity=None, attrs=None, ignore=None, params=None):
+        """Provide a default value for ``entity``.
+
+        By default, ``nailgun.entity_mixins.EntityReadMixin.read`` provides a
+        default value for ``entity`` like so::
+
+            entity = type(self)()
+
+        However, :class:`SSHKey` requires that an ``user`` be
+        provided, so this technique will not work. Do this instead::
+
+            entity = type(self)(user=self.user.id)
+
+        """
+        # read() should not change the state of the object it's called on, but
+        # super() alters the attributes of any entity passed in. Creating a new
+        # object and passing it to super() lets this one avoid changing state.
+        if entity is None:
+            entity = type(self)(
+                self._server_config,
+                user=self.user,  # pylint:disable=no-member
+            )
+        if ignore is None:
+            ignore = set()
+        ignore.add('user')
+        return super(SSHKey, self).read(entity, attrs, ignore, params)
+
+    def search_normalize(self, results):
+        """Append user id to search results to be able to initialize found
+        :class:`User` successfully
+        """
+        for sshkey in results:
+            sshkey[u'user_id'] = self.user.id  # pylint:disable=no-member
+        return super(SSHKey, self).search_normalize(results)
+
+
 class Status(Entity):
     """A representation of a Status entity."""
 
diff --git a/tests/test_entities.py b/tests/test_entities.py
index 298f4ff0..cc925d35 100644
--- a/tests/test_entities.py
+++ b/tests/test_entities.py
@@ -197,6 +197,7 @@ def test_init_succeeds(self):
             (entities.Parameter, {'organization': 1}),
             (entities.Parameter, {'subnet': 1}),
             (entities.RepositorySet, {'product': 1}),
+            (entities.SSHKey, {'user': 1}),
             (entities.SyncPlan, {'organization': 1}),
         ])
         for entity, params in entities_:
@@ -1030,6 +1031,7 @@ def test_entity_arg(self):
                 entities.Parameter(self.cfg, organization=2),
                 entities.Parameter(self.cfg, subnet=2),
                 entities.RepositorySet(self.cfg, product=2),
+                entities.SSHKey(self.cfg, user=2),
                 entities.SyncPlan(self.cfg, organization=2),
         ):
             # We mock read_json() because it may be called by read().
@@ -1363,6 +1365,25 @@ def setUpClass(cls):
         """Set a server configuration at ``cls.cfg``."""
         cls.cfg = config.ServerConfig('http://example.com')
 
+    def test_sshkey(self):
+        """Test :meth:`nailgun.entities.SSHKey.search_normalize`.
+
+        Assert that ``user_id`` was added with correct user's id to search
+        results.
+        """
+        results = [
+            {'id': 1, 'login': 'foo'},
+            {'id': 2, 'login': 'bar'},
+        ]
+        with mock.patch.object(
+            EntitySearchMixin,
+            'search_normalize',
+        ) as search_normalize:
+            entities.SSHKey(self.cfg, user=4).search_normalize(results)
+            for args in search_normalize.call_args[0][0]:
+                self.assertIn('user_id', args)
+                self.assertEqual(args['user_id'], 4)
+
     def test_interface(self):
         """Test :meth:`nailgun.entities.Interface.search_normalize`.