From ced24273ce8d191a95b8222d278e9725476e6e27 Mon Sep 17 00:00:00 2001
From: Andrei <b1naryth1ef@gmail.com>
Date: Tue, 2 Jul 2019 14:30:43 -0700
Subject: [PATCH] ref: improve User.get_avatar_url and Guild.get_icon_url

---
 disco/types/guild.py | 13 ++++++++-----
 disco/types/user.py  | 13 ++++++++-----
 tests/types/user.py  | 33 ++++++++++++++++-----------------
 3 files changed, 32 insertions(+), 27 deletions(-)

diff --git a/disco/types/guild.py b/disco/types/guild.py
index d9ca2e5..db49430 100644
--- a/disco/types/guild.py
+++ b/disco/types/guild.py
@@ -526,15 +526,18 @@ class Guild(SlottedModel, Permissible):
     def get_emoji(self, emoji):
         return self.client.api.guilds_emojis_get(self.id, emoji)
 
-    def get_icon_url(self, fmt=None, size=1024):
+    def get_icon_url(self, still_format='webp', animated_format='gif', size=1024):
         if not self.icon:
             return ''
-        if fmt is not None:
-            return 'https://cdn.discordapp.com/icons/{}/{}.{}?size={}'.format(self.id, self.icon, fmt, size)
+
         if self.icon.startswith('a_'):
-            return 'https://cdn.discordapp.com/icons/{}/{}.gif?size={}'.format(self.id, self.icon, size)
+            return 'https://cdn.discordapp.com/icons/{}/{}.{}?size={}'.format(
+                self.id, self.icon, animated_format, size
+            )
         else:
-            return 'https://cdn.discordapp.com/icons/{}/{}.webp?size={}'.format(self.id, self.icon, size)
+            return 'https://cdn.discordapp.com/icons/{}/{}.{}?size={}'.format(
+                self.id, self.icon, still_format, size
+            )
 
     def get_splash_url(self, fmt='webp', size=1024):
         if not self.splash:
diff --git a/disco/types/user.py b/disco/types/user.py
index 2791427..30c432c 100644
--- a/disco/types/user.py
+++ b/disco/types/user.py
@@ -27,15 +27,18 @@ class User(SlottedModel, with_equality('id'), with_hash('id')):
 
     presence = Field(None)
 
-    def get_avatar_url(self, fmt=None, size=1024):
+    def get_avatar_url(self, still_format='webp', animated_format='gif', size=1024):
         if not self.avatar:
             return 'https://cdn.discordapp.com/embed/avatars/{}.png'.format(self.default_avatar)
-        if fmt is not None:
-            return 'https://cdn.discordapp.com/avatars/{}/{}.{}?size={}'.format(self.id, self.avatar, fmt, size)
+
         if self.avatar.startswith('a_'):
-            return 'https://cdn.discordapp.com/avatars/{}/{}.gif?size={}'.format(self.id, self.avatar, size)
+            return 'https://cdn.discordapp.com/avatars/{}/{}.{}?size={}'.format(
+                self.id, self.avatar, animated_format, size
+            )
         else:
-            return 'https://cdn.discordapp.com/avatars/{}/{}.webp?size={}'.format(self.id, self.avatar, size)
+            return 'https://cdn.discordapp.com/avatars/{}/{}.{}?size={}'.format(
+                self.id, self.avatar, still_format, size
+            )
 
     @property
     def default_avatar(self):
diff --git a/tests/types/user.py b/tests/types/user.py
index 89eaca6..f517d9e 100644
--- a/tests/types/user.py
+++ b/tests/types/user.py
@@ -1,22 +1,21 @@
-from unittest import TestCase
-
 from disco.types.user import User, DefaultAvatars
 
 
-class TestChannel(TestCase):
-    def test_user_avatar(self):
-        u = User(
-            id=12345,
-            username='test123',
-            avatar='1234567890abcdefghijkl',
-            discriminator='1234',
-            bot=False)
+def test_user_avatar_url():
+    u = User(id=12345, avatar='1234567890abcdefghijkl')
+    assert u.avatar_url == 'https://cdn.discordapp.com/avatars/12345/1234567890abcdefghijkl.webp?size=1024'
+    avatar_url = u.get_avatar_url(still_format='png')
+    assert avatar_url == 'https://cdn.discordapp.com/avatars/12345/1234567890abcdefghijkl.png?size=1024'
+
+
+def test_user_animated_avatar_url():
+    u = User(id=12345, avatar='a_1234567890abcdefghijkl')
+    assert u.avatar_url == 'https://cdn.discordapp.com/avatars/12345/a_1234567890abcdefghijkl.gif?size=1024'
+    avatar_url = u.get_avatar_url(animated_format='webp')
+    assert avatar_url == 'https://cdn.discordapp.com/avatars/12345/a_1234567890abcdefghijkl.webp?size=1024'
 
-        self.assertEqual(
-            u.avatar_url, 'https://cdn.discordapp.com/avatars/12345/1234567890abcdefghijkl.webp?size=1024'
-        )
 
-    def test_user_default_avatar(self):
-        u = User(id=123456, discriminator='1234')
-        self.assertEqual(u.default_avatar, DefaultAvatars.RED)
-        self.assertEqual(u.avatar_url, 'https://cdn.discordapp.com/embed/avatars/4.png')
+def test_user_default_avatar_url():
+    u = User(id=12345, discriminator='1234')
+    assert u.default_avatar == DefaultAvatars.RED
+    assert u.avatar_url == 'https://cdn.discordapp.com/embed/avatars/4.png'