Add a SetMixin

Gauvain Pocentek · Gauvain Pocentek · commit a1b097ce1811 · 2017-11-11T15:40:13.000+01:00
Use it for UserCustomAttribute, will be useful for {Project,Group}CustomAttribute (python-gitlab#367)
diff --git a/gitlab/mixins.py b/gitlab/mixins.py
@@ -222,6 +222,29 @@ def update(self, id=None, new_data={}, **kwargs):
         return self.gitlab.http_put(path, post_data=data, **kwargs)
 
 
+class SetMixin(object):
+    @exc.on_http_error(exc.GitlabSetError)
+    def set(self, key, value, **kwargs):
+        """Create or update the object.
+
+        Args:
+            key (str): The key of the object to create/update
+            value (str): The value to set for the object
+            **kwargs: Extra options to send to the server (e.g. sudo)
+
+        Raises:
+            GitlabAuthenticationError: If authentication is not correct
+            GitlabSetError: If an error occured
+
+        Returns:
+            UserCustomAttribute: The created/updated user attribute
+        """
+        path = '%s/%s' % (self.path, key.replace('/', '%2F'))
+        data = {'value': value}
+        server_data = self.gitlab.http_put(path, post_data=data, **kwargs)
+        return self._obj_cls(self, server_data)
+
+
 class DeleteMixin(object):
     @exc.on_http_error(exc.GitlabDeleteError)
     def delete(self, id, **kwargs):
diff --git a/gitlab/tests/test_mixins.py b/gitlab/tests/test_mixins.py
@@ -66,6 +66,13 @@ class O(TimeTrackingMixin):
         self.assertTrue(hasattr(obj, 'add_spent_time'))
         self.assertTrue(hasattr(obj, 'reset_spent_time'))
 
+    def test_set_mixin(self):
+        class O(SetMixin):
+            pass
+
+        obj = O()
+        self.assertTrue(hasattr(obj, 'set'))
+
 
 class TestMetaMixins(unittest.TestCase):
     def test_retrieve_mixin(self):
@@ -409,3 +416,21 @@ def resp_cont(url, request):
             obj.save()
             self.assertEqual(obj._attrs['foo'], 'baz')
             self.assertDictEqual(obj._updated_attrs, {})
+
+    def test_set_mixin(self):
+        class M(SetMixin, FakeManager):
+            pass
+
+        @urlmatch(scheme="http", netloc="localhost", path='/api/v4/tests/foo',
+                  method="put")
+        def resp_cont(url, request):
+            headers = {'Content-Type': 'application/json'}
+            content = '{"key": "foo", "value": "bar"}'
+            return response(200, content, headers, None, 5, request)
+
+        with HTTMock(resp_cont):
+            mgr = M(self.gl)
+            obj = mgr.set('foo', 'bar')
+            self.assertIsInstance(obj, FakeObject)
+            self.assertEqual(obj.key, 'foo')
+            self.assertEqual(obj.value, 'bar')
diff --git a/gitlab/v4/objects.py b/gitlab/v4/objects.py
@@ -125,31 +125,12 @@ class UserCustomAttribute(ObjectDeleteMixin, RESTObject):
     _id_attr = 'key'
 
 
-class UserCustomAttributeManager(RetrieveMixin, DeleteMixin, RESTManager):
+class UserCustomAttributeManager(RetrieveMixin, SetMixin, DeleteMixin,
+                                 RESTManager):
     _path = '/users/%(user_id)s/custom_attributes'
     _obj_cls = UserCustomAttribute
     _from_parent_attrs = {'user_id': 'id'}
 
-    def set(self, key, value, **kwargs):
-        """Create or update a user attribute.
-
-        Args:
-            key (str): The attribute to update
-            value (str): The value to set
-            **kwargs: Extra options to send to the server (e.g. sudo)
-
-        Raises:
-            GitlabAuthenticationError: If authentication is not correct
-            GitlabSetError: If an error occured
-
-        Returns:
-            UserCustomAttribute: The created/updated user attribute
-        """
-        path = '%s/%s' % (self.path, key.replace('/', '%2F'))
-        data = {'value': value}
-        server_data = self.gitlab.http_put(path, post_data=data, **kwargs)
-        return self._obj_cls(self, server_data)
-
 
 class UserEmail(ObjectDeleteMixin, RESTObject):
     _short_print_attr = 'email'
