Add support for searching genre stations. (#45)

Add support for searching genre stations.

3 files changed, 223 insertions, 47 deletions

diff --git a/pandora/client.py b/pandora/client.py
index 7c9791e..48077cb 100644
--- a/pandora/client.py
+++ b/pandora/client.py
@@ -174,12 +174,18 @@ class APIClient(BaseAPIClient):
         return self("bookmark.deleteArtistBookmark",
                     bookmarkToken=bookmark_token)
 
-    def search(self, search_text):
+    def search(self, search_text,
+               include_near_matches=False,
+               include_genre_stations=False):
         from .models.pandora import SearchResult
 
-        return SearchResult.from_json(self,
-                                      self("music.search",
-                                           searchText=search_text))
+        return SearchResult.from_json(
+            self,
+            self("music.search",
+                 searchText=search_text,
+                 includeNearMatches=include_near_matches,
+                 includeGenreStations=include_genre_stations)
+        )
 
     def add_feedback(self, track_token, positive):
         return self("station.addFeedback",
diff --git a/pandora/models/pandora.py b/pandora/models/pandora.py
index f2fd36e..8fa8e21 100644
--- a/pandora/models/pandora.py
+++ b/pandora/models/pandora.py
@@ -284,37 +284,98 @@ class BookmarkList(PandoraModel):
 
 class SearchResultItem(PandoraModel):
 
-    artist = Field("artistName")
-    song_name = Field("songName")
     score = Field("score")
-    likely_match = Field("likelyMatch", default=False)
     token = Field("musicToken")
 
     @property
     def is_song(self):
-        return self.token.startswith('S')
+        return isinstance(self, SongSearchResultItem)
 
     @property
     def is_artist(self):
-        return self.token.startswith('R')
+        return isinstance(self, ArtistSearchResultItem) and \
+               self.token.startswith("R")
 
     @property
     def is_composer(self):
-        return self.token.startswith('C')
+        return isinstance(self, ArtistSearchResultItem) and \
+               self.token.startswith("C")
+
+    @property
+    def is_genre_station(self):
+        return isinstance(self, GenreStationSearchResultItem)
 
     def create_station(self):
-        if self.is_song:
-            self._api_client.create_station(track_token=self.token)
+        raise NotImplementedError
+
+    @classmethod
+    def from_json(cls, api_client, data):
+        if data["musicToken"].startswith("S"):
+            return SongSearchResultItem.from_json(api_client, data)
+
+        elif data["musicToken"].startswith(("R", "C")):
+            return ArtistSearchResultItem.from_json(api_client, data)
+
+        elif data["musicToken"].startswith("G"):
+            return GenreStationSearchResultItem.from_json(api_client, data)
         else:
-            self._api_client.create_station(artist_token=self.token)
+            raise NotImplementedError("Unknown result token type '{}'"
+                                      .format(data["musicToken"]))
+
+
+class ArtistSearchResultItem(SearchResultItem):
+
+    score = Field("score")
+    token = Field("musicToken")
+    artist = Field("artistName")
+    likely_match = Field("likelyMatch", default=False)
+
+    def create_station(self):
+        self._api_client.create_station(artist_token=self.token)
+
+    @classmethod
+    def from_json(cls, api_client, data):
+        return super(SearchResultItem, cls).from_json(api_client, data)
+
+
+class SongSearchResultItem(SearchResultItem):
+
+    score = Field("score")
+    token = Field("musicToken")
+    artist = Field("artistName")
+    song_name = Field("songName")
+
+    def create_station(self):
+        self._api_client.create_station(track_token=self.token)
+
+    @classmethod
+    def from_json(cls, api_client, data):
+        return super(SearchResultItem, cls).from_json(api_client, data)
+
+
+class GenreStationSearchResultItem(SearchResultItem):
+
+    score = Field("score")
+    token = Field("musicToken")
+    station_name = Field("stationName")
+
+    def create_station(self):
+        self._api_client.create_station(search_token=self.token)
+
+    @classmethod
+    def from_json(cls, api_client, data):
+        return super(SearchResultItem, cls).from_json(api_client, data)
 
 
 class SearchResult(PandoraModel):
 
     nearest_matches_available = Field("nearMatchesAvailable")
     explanation = Field("explanation")
-    songs = Field("songs", formatter=SearchResultItem.from_json_list)
-    artists = Field("artists", formatter=SearchResultItem.from_json_list)
+    songs = Field("songs", formatter=SongSearchResultItem.from_json_list)
+    artists = Field("artists", formatter=ArtistSearchResultItem.from_json_list)
+    genre_stations = Field(
+        "genreStations",
+        formatter=GenreStationSearchResultItem.from_json_list)
 
 
 class GenreStationList(PandoraDictListModel):
diff --git a/tests/test_pandora/test_models.py b/tests/test_pandora/test_models.py
index 7383d85..3a771ab 100644
--- a/tests/test_pandora/test_models.py
+++ b/tests/test_pandora/test_models.py
@@ -277,54 +277,157 @@ class TestAdItem(TestCase):
 
 class TestSearchResultItem(TestCase):
 
-    JSON_DATA = {
+    SONG_JSON_DATA = {
         "artistName": "artist_name_mock",
         "musicToken": "S0000000",
         "songName": "song_name_mock",
         "score": 100
     }
 
+    ARTIST_JSON_DATA = {
+        "artistName": "artist_name_mock",
+        "musicToken": "R0000000",
+        "likelyMatch": False,
+        "score": 100
+    }
+
+    COMPOSER_JSON_DATA = {
+        "artistName": "composer_name_mock",
+        "musicToken": "C0000000",
+        "likelyMatch": False,
+        "score": 100
+    }
+
+    GENRE_JSON_DATA = {
+        "stationName": "station_name_mock",
+        "musicToken": "G0000000",
+        "score": 100
+    }
+
     def setUp(self):
-        api_client_mock = Mock(spec=APIClient)
-        api_client_mock.default_audio_quality = APIClient.HIGH_AUDIO_QUALITY
-        self.result = SearchResultItem.from_json(api_client_mock, self.JSON_DATA)
+        self.api_client_mock = Mock(spec=APIClient)
+        self.api_client_mock.default_audio_quality = APIClient.HIGH_AUDIO_QUALITY
+
+    def test_is_song(self):
+        result = SearchResultItem.from_json(self.api_client_mock, self.SONG_JSON_DATA)
+        assert result.is_song
+        assert not result.is_artist
+        assert not result.is_composer
+        assert not result.is_genre_station
+
+    def test_is_artist(self):
+        result = SearchResultItem.from_json(self.api_client_mock, self.ARTIST_JSON_DATA)
+        assert not result.is_song
+        assert result.is_artist
+        assert not result.is_composer
+        assert not result.is_genre_station
+
+    def test_is_composer(self):
+        result = SearchResultItem.from_json(self.api_client_mock, self.COMPOSER_JSON_DATA)
+        assert not result.is_song
+        assert not result.is_artist
+        assert result.is_composer
+        assert not result.is_genre_station
+
+    def test_is_genre_station(self):
+        result = SearchResultItem.from_json(self.api_client_mock, self.GENRE_JSON_DATA)
+        assert not result.is_song
+        assert not result.is_artist
+        assert not result.is_composer
+        assert result.is_genre_station
+
+    def test_create_station(self):
+        result = SearchResultItem.from_json(self.api_client_mock, self.SONG_JSON_DATA)
+
+        self.assertRaises(NotImplementedError, result.create_station())
+
+
+class TestArtistSearchResultItem(TestCase):
+
+    ARTIST_JSON_DATA = {
+        "artistName": "artist_name_mock",
+        "musicToken": "R0000000",
+        "likelyMatch": False,
+        "score": 100
+    }
+
+    COMPOSER_JSON_DATA = {
+        "artistName": "composer_name_mock",
+        "musicToken": "C0000000",
+        "likelyMatch": False,
+        "score": 100
+    }
+
+    def setUp(self):
+        self.api_client_mock = Mock(spec=APIClient)
+        self.api_client_mock.default_audio_quality = APIClient.HIGH_AUDIO_QUALITY
 
     def test_repr(self):
-        expected = ("SearchResultItem(artist='artist_name_mock', likely_match=False, score=100, "
-                    "song_name='song_name_mock', token='S0000000')")
-        self.assertEqual(expected, repr(self.result))
+        result = SearchResultItem.from_json(self.api_client_mock, self.ARTIST_JSON_DATA)
+        expected = ("ArtistSearchResultItem(artist='artist_name_mock', likely_match=False, score=100, token='R0000000')")
+        self.assertEqual(expected, repr(result))
 
-    def test_is_song(self):
-        assert self.result.is_song
+        result = SearchResultItem.from_json(self.api_client_mock, self.COMPOSER_JSON_DATA)
+        expected = ("ArtistSearchResultItem(artist='composer_name_mock', likely_match=False, score=100, token='C0000000')")
+        self.assertEqual(expected, repr(result))
 
-        self.result.token = 'R123456'
-        assert not self.result.is_song
+    def test_create_station(self):
+        result = SearchResultItem.from_json(self.api_client_mock, self.ARTIST_JSON_DATA)
+        result._api_client.create_station = Mock()
 
-    def test_is_artist(self):
-        assert not self.result.is_artist
+        result.create_station()
+        result._api_client.create_station.assert_called_with(artist_token=result.token)
 
-        self.result.token = 'R123456'
-        assert self.result.is_artist
 
-    def test_is_composer(self):
-        assert not self.result.is_composer
+class TestSongSearchResultItem(TestCase):
 
-        self.result.token = 'C12345'
-        assert self.result.is_composer
+    SONG_JSON_DATA = {
+        "artistName": "artist_name_mock",
+        "musicToken": "S0000000",
+        "songName": "song_name_mock",
+        "score": 100
+    }
 
+    def setUp(self):
+        self.api_client_mock = Mock(spec=APIClient)
+        self.api_client_mock.default_audio_quality = APIClient.HIGH_AUDIO_QUALITY
 
-    def test_create_station_song(self):
-        self.result._api_client.create_station = Mock()
+    def test_repr(self):
+        result = SearchResultItem.from_json(self.api_client_mock, self.SONG_JSON_DATA)
+        expected = ("SongSearchResultItem(artist='artist_name_mock', score=100, song_name='song_name_mock', token='S0000000')")
+        self.assertEqual(expected, repr(result))
+
+    def test_create_station(self):
+        result = SearchResultItem.from_json(self.api_client_mock, self.SONG_JSON_DATA)
+        result._api_client.create_station = Mock()
+
+        result.create_station()
+        result._api_client.create_station.assert_called_with(track_token=result.token)
+
+
+class TestGenreStationSearchResultItem(TestCase):
+
+    GENRE_JSON_DATA = {
+        "stationName": "station_name_mock",
+        "musicToken": "G0000000",
+        "score": 100
+    }
 
-        self.result.create_station()
-        self.result._api_client.create_station.assert_called_with(track_token=self.result.token)
+    def setUp(self):
+        self.api_client_mock = Mock(spec=APIClient)
+        self.api_client_mock.default_audio_quality = APIClient.HIGH_AUDIO_QUALITY
 
-    def test_create_station_artist(self):
-        self.result.token = 'R123456'
-        self.result._api_client.create_station = Mock()
+    def test_repr(self):
+        result = SearchResultItem.from_json(self.api_client_mock, self.GENRE_JSON_DATA)
+        expected = ("GenreStationSearchResultItem(score=100, station_name='station_name_mock', token='G0000000')")
+        self.assertEqual(expected, repr(result))
+
+    def test_create_station(self):
+        result = SearchResultItem.from_json(self.api_client_mock, self.GENRE_JSON_DATA)
+        result._api_client.create_station = Mock()
 
-        self.result.create_station()
-        self.result._api_client.create_station.assert_called_with(artist_token=self.result.token)
+        result.create_station()
+        result._api_client.create_station.assert_called_with(search_token=result.token)
 
 
 class TestSearchResult(TestCase):
@@ -343,6 +446,11 @@ class TestSearchResult(TestCase):
             'musicToken': 'R000000',
             'likelyMatch': False,
             'score': 80
+        }],
+        'genreStations': [{
+            'musicToken': 'G0000',
+            'stationName': 'station_mock',
+            'score': 50
         }]
     }
 
@@ -352,8 +460,9 @@ class TestSearchResult(TestCase):
         self.result = SearchResult.from_json(api_client_mock, self.JSON_DATA)
 
     def test_repr(self):
-        expected = ("SearchResult(artists=[SearchResultItem(artist='artist_mock', "
-                    "likely_match=False, score=80, song_name=None, token='R000000')], explanation='', "
-                    "nearest_matches_available=True, songs=[SearchResultItem(artist='song_artist_mock', "
-                    "likely_match=False, score=100, song_name='song_name_mock', token='S0000000')])")
+        expected = ("SearchResult(artists=[ArtistSearchResultItem(artist='artist_mock', likely_match=False, score=80, "
+                    "token='R000000')], explanation='', genre_stations=[GenreStationSearchResultItem(score=50, "
+                    "station_name='station_mock', token='G0000')], nearest_matches_available=True, "
+                    "songs=[SongSearchResultItem(artist='song_artist_mock', score=100, song_name='song_name_mock', "
+                    "token='S0000000')])")
         self.assertEqual(expected, repr(self.result))