Enhancement: Adds additional audio urls (#58)

* Adds support for additionalAudioUrl for station.getPlaylist
* Fixes broken tests
* Reworks use of iterables for additional audio urls and adds associated tests
* Moves parsing additional url response into a syntethic field
* Adds tests for additional urls field

4 files changed, 158 insertions, 10 deletions

diff --git a/pandora/client.py b/pandora/client.py
index 01b7654..df09144 100644
--- a/pandora/client.py
+++ b/pandora/client.py
@@ -127,15 +127,28 @@ class APIClient(BaseAPIClient):
     def get_station_list_checksum(self):  # pragma: no cover
         return self("user.getStationListChecksum")["checksum"]
 
-    def get_playlist(self, station_token):
+    def get_playlist(self, station_token, additional_urls=None):
         from .models.pandora import Playlist
 
-        playlist = Playlist.from_json(self,
-                                      self("station.getPlaylist",
-                                           stationToken=station_token,
-                                           includeTrackLength=True,
-                                           xplatformAdCapable=True,
-                                           audioAdPodCapable=True))
+        if additional_urls is None:
+            additional_urls = []
+
+        if isinstance(additional_urls, str):
+            raise TypeError('Additional urls should be a list')
+
+        urls = [getattr(url, "value", url) for url in additional_urls]
+
+        resp = self("station.getPlaylist",
+                    stationToken=station_token,
+                    includeTrackLength=True,
+                    xplatformAdCapable=True,
+                    audioAdPodCapable=True,
+                    additionalAudioUrl=','.join(urls))
+
+        for item in resp['items']:
+            item['_paramAdditionalUrls'] = additional_urls
+
+        playlist = Playlist.from_json(self, resp)
 
         for i, track in enumerate(playlist):
             if track.is_ad:
diff --git a/pandora/models/pandora.py b/pandora/models/pandora.py
index cc57eac..ecc1225 100644
--- a/pandora/models/pandora.py
+++ b/pandora/models/pandora.py
@@ -6,6 +6,18 @@ from . import Field, DateField, SyntheticField
 from . import PandoraModel, PandoraListModel, PandoraDictListModel
 
 
+class AdditionalAudioUrl(Enum):
+    HTTP_40_AAC_MONO = 'HTTP_40_AAC_MONO'
+    HTTP_64_AAC = 'HTTP_64_AAC'
+    HTTP_32_AACPLUS = 'HTTP_32_AACPLUS'
+    HTTP_64_AACPLUS = 'HTTP_64_AACPLUS'
+    HTTP_24_AACPLUS_ADTS = 'HTTP_24_AACPLUS_ADTS'
+    HTTP_32_AACPLUS_ADTS = 'HTTP_32_AACPLUS_ADTS'
+    HTTP_64_AACPLUS_ADTS = 'HTTP_64_AACPLUS_ADTS'
+    HTTP_128_MP3 = 'HTTP_128_MP3'
+    HTTP_32_WMA = 'HTTP_32_WMA'
+
+
 class PandoraType(Enum):
 
     TRACK = "TR"
@@ -100,8 +112,9 @@ class Station(PandoraModel):
     seeds = Field("music", model=StationSeeds)
     feedback = Field("feedback", model=StationFeedback)
 
-    def get_playlist(self):
-        return iter(self._api_client.get_playlist(self.token))
+    def get_playlist(self, additional_urls=None):
+        return iter(self._api_client.get_playlist(self.token,
+                                                  additional_urls))
 
 
 class GenreStation(PandoraModel):
@@ -178,6 +191,27 @@ class AudioField(SyntheticField):
         return audio_url[self.field] if audio_url else None
 
 
+class AdditionalUrlField(SyntheticField):
+
+    def formatter(self, api_client, data, value):
+        """Parse additional url fields and map them to inputs
+
+        Attempt to create a dictionary with keys being user input, and
+        response being the returned URL
+        """
+        if value is None:
+            return None
+
+        user_param = data['_paramAdditionalUrls']
+        urls = {}
+        if isinstance(value, str):
+            urls[user_param[0]] = value
+        else:
+            for key, url in zip(user_param, value):
+                urls[key] = url
+        return urls
+
+
 class PlaylistModel(PandoraModel):
 
     def get_is_playable(self):
@@ -243,6 +277,8 @@ class PlaylistItem(PlaylistModel):
     song_detail_url = Field("songDetailUrl")
     song_explore_url = Field("songExplorerUrl")
 
+    additional_audio_urls = AdditionalUrlField("additionalAudioUrl")
+
     @property
     def is_ad(self):
         return self.ad_token is not None
diff --git a/tests/test_pandora/test_client.py b/tests/test_pandora/test_client.py
index 6c61355..d24ea28 100644
--- a/tests/test_pandora/test_client.py
+++ b/tests/test_pandora/test_client.py
@@ -1,7 +1,7 @@
 from unittest import TestCase
 
 from pandora import errors
-from pandora.models.pandora import AdItem
+from pandora.models.pandora import AdItem, AdditionalAudioUrl
 from pandora.client import APIClient, BaseAPIClient
 from pandora.py2compat import Mock, call, patch
 from tests.test_pandora.test_models import TestAdItem
@@ -87,6 +87,7 @@ class TestCallingAPIClient(TestCase):
             client.get_playlist('token_mock')
 
             playlist_mock.assert_has_calls([call("station.getPlaylist",
+                                                 additionalAudioUrl='',
                                                  audioAdPodCapable=True,
                                                  includeTrackLength=True,
                                                  stationToken='token_mock',
@@ -186,3 +187,68 @@ class TestCreatingGenreStation(TestCase):
             client = APIClient(Mock(), None, None, None, None)
             station = client.get_genre_stations()
             self.assertEqual(station.checksum, "foo")
+
+
+class TestAdditionalUrls(TestCase):
+
+    def test_non_iterable_string(self):
+        with self.assertRaises(TypeError):
+            transport = Mock(side_effect=[errors.InvalidAuthToken(), None])
+
+            client = APIClient(transport, None, None, None, None)
+            client._authenticate = Mock()
+
+            client.get_playlist('token_mock', additional_urls='')
+
+    def test_non_iterable_other(self):
+        with self.assertRaises(TypeError):
+            transport = Mock(side_effect=[errors.InvalidAuthToken(), None])
+
+            client = APIClient(transport, None, None, None, None)
+            client._authenticate = Mock()
+
+            client.get_playlist('token_mock',
+                                additional_urls=AdditionalAudioUrl.HTTP_32_WMA)
+
+    def test_without_enum(self):
+        with patch.object(APIClient, '__call__') as playlist_mock:
+            transport = Mock(side_effect=[errors.InvalidAuthToken(), None])
+
+            client = APIClient(transport, None, None, None, None)
+            client._authenticate = Mock()
+
+            urls = ['HTTP_128_MP3',
+                    'HTTP_24_AACPLUS_ADTS']
+
+            desired = 'HTTP_128_MP3,HTTP_24_AACPLUS_ADTS'
+
+            client.get_playlist('token_mock', additional_urls=urls)
+
+            playlist_mock.assert_has_calls([call("station.getPlaylist",
+                                                 additionalAudioUrl=desired,
+                                                 audioAdPodCapable=True,
+                                                 includeTrackLength=True,
+                                                 stationToken='token_mock',
+                                                 xplatformAdCapable=True)])
+
+
+    def test_with_enum(self):
+        with patch.object(APIClient, '__call__') as playlist_mock:
+            transport = Mock(side_effect=[errors.InvalidAuthToken(), None])
+
+            client = APIClient(transport, None, None, None, None)
+            client._authenticate = Mock()
+
+            urls = [AdditionalAudioUrl.HTTP_128_MP3,
+                    AdditionalAudioUrl.HTTP_24_AACPLUS_ADTS]
+
+            desired = 'HTTP_128_MP3,HTTP_24_AACPLUS_ADTS'
+
+            client.get_playlist('token_mock', additional_urls=urls)
+
+            playlist_mock.assert_has_calls([call("station.getPlaylist",
+                                                 additionalAudioUrl=desired,
+                                                 audioAdPodCapable=True,
+                                                 includeTrackLength=True,
+                                                 stationToken='token_mock',
+                                                 xplatformAdCapable=True)])
diff --git a/tests/test_pandora/test_models.py b/tests/test_pandora/test_models.py
index ce71ef8..410b0de 100644
--- a/tests/test_pandora/test_models.py
+++ b/tests/test_pandora/test_models.py
@@ -52,6 +52,39 @@ class TestDateField(TestCase):
         self.assertEqual(expected, model.date_field.replace(microsecond=0))
 
 
+class TestAdditionalUrlField(TestCase):
+
+    def test_single_url(self):
+        dummy_data = {
+                '_paramAdditionalUrls': ['foo']
+                }
+
+        field = pm.AdditionalUrlField("additionalAudioUrl")
+
+        ret = field.formatter(None, dummy_data, 'test')
+
+        self.assertEqual(ret, {'foo': 'test'})
+
+    def test_multiple_urls(self):
+        dummy_data = {
+                '_paramAdditionalUrls': [
+                    'abc',
+                    'def',
+                    ]
+                }
+
+        field = pm.AdditionalUrlField("additionalAudioUrl")
+
+        ret = field.formatter(None, dummy_data, ['foo', 'bar'])
+
+        expected = {
+                'abc': 'foo',
+                'def': 'bar',
+                }
+
+        self.assertEqual(ret, expected)
+
+
 class TestPandoraModel(TestCase):
 
     JSON_DATA = {