Add type to User and GroupChat classes

leandrotoledo · leandrotoledo · commit cf5d18476673 · 2015-10-08T10:19:05.000-03:00
diff --git a/telegram/groupchat.py b/telegram/groupchat.py
@@ -28,18 +28,22 @@ class GroupChat(TelegramObject):
     Attributes:
         id (int):
         title (str):
+        type (str):
 
     Args:
         id (int):
         title (str):
+        type (str):
     """
 
     def __init__(self,
                  id,
-                 title):
+                 title,
+                 type):
         # Required
         self.id = int(id)
         self.title = title
+        self.type = type
 
     @staticmethod
     def de_json(data):
diff --git a/telegram/user.py b/telegram/user.py
@@ -30,13 +30,15 @@ class User(TelegramObject):
         first_name (str):
         last_name (str):
         username (str):
+        type (str):
 
     Args:
         id (int):
         first_name (str):
         **kwargs: Arbitrary keyword arguments.
 
     Keyword Args:
+        type (Optional[str]):
         last_name (Optional[str]):
         username (Optional[str]):
     """
@@ -49,6 +51,7 @@ def __init__(self,
         self.id = int(id)
         self.first_name = first_name
         # Optionals
+        self.type = kwargs.get('type', '')
         self.last_name = kwargs.get('last_name', '')
         self.username = kwargs.get('username', '')
 
diff --git a/tests/test_groupchat.py b/tests/test_groupchat.py
@@ -33,10 +33,12 @@ class GroupChatTest(BaseTest, unittest.TestCase):
     def setUp(self):
         self.id = -28767330
         self.title = 'ToledosPalaceBot - Group'
+        self.type = 'group'
 
         self.json_dict = {
             'id': self.id,
-            'title': self.title
+            'title': self.title,
+            'type': self.type
         }
 
     def test_group_chat_de_json_empty_json(self):
@@ -55,6 +57,7 @@ def test_group_chat_de_json(self):
 
         self.assertEqual(group_chat.id, self.id)
         self.assertEqual(group_chat.title, self.title)
+        self.assertEqual(group_chat.type, self.type)
 
     def test_group_chat_to_json(self):
         """Test GroupChat.to_json() method"""
@@ -73,6 +76,7 @@ def test_group_chat_to_dict(self):
         self.assertTrue(self.is_dict(group_chat.to_dict()))
         self.assertEqual(group_chat['id'], self.id)
         self.assertEqual(group_chat['title'], self.title)
+        self.assertEqual(group_chat['type'], self.type)
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_user.py b/tests/test_user.py
@@ -35,12 +35,14 @@ def setUp(self):
         self.first_name = "Leandro"
         self.last_name = "S."
         self.username = "leandrotoledo"
+        self.type = "private"
 
         self.json_dict = {
             'id': self.id,
             'first_name': self.first_name,
             'last_name': self.last_name,
-            'username': self.username
+            'username': self.username,
+            'type': self.type
         }
 
     def test_user_de_json(self):
@@ -53,6 +55,7 @@ def test_user_de_json(self):
         self.assertEqual(user.first_name, self.first_name)
         self.assertEqual(user.last_name, self.last_name)
         self.assertEqual(user.username, self.username)
+        self.assertEqual(user.type, self.type)
 
         self.assertEqual(user.name, '@leandrotoledo')
 
@@ -69,6 +72,7 @@ def test_user_de_json_without_username(self):
         self.assertEqual(user.id, self.id)
         self.assertEqual(user.first_name, self.first_name)
         self.assertEqual(user.last_name, self.last_name)
+        self.assertEqual(user.type, self.type)
 
         self.assertEqual(user.name, '%s %s' % (self.first_name, self.last_name))
 
@@ -108,6 +112,7 @@ def test_user_to_dict(self):
         self.assertEqual(user['first_name'], self.first_name)
         self.assertEqual(user['last_name'], self.last_name)
         self.assertEqual(user['username'], self.username)
+        self.assertEqual(user['type'], self.type)
 
 if __name__ == '__main__':
     unittest.main()
