diff --git a/src/main/java/musicplayer/PlayerGUI.java b/src/main/java/musicplayer/PlayerGUI.java
index 4d51bbe..3d4fbb4 100644
--- a/src/main/java/musicplayer/PlayerGUI.java
+++ b/src/main/java/musicplayer/PlayerGUI.java
@@ -8,6 +8,8 @@ import musicplayer.callbacks.PlayerCallbackInterface;
 import musicplayer.db.DatabaseManager;
 import musicplayer.db.Gateway;
 import musicplayer.model.Album;
+import musicplayer.model.Artist;
+import musicplayer.model.HasSongs;
 import musicplayer.model.Song;
 import musicplayer.swingmodels.PlaylistTableModel;
 import musicplayer.util.LibraryUtils;
@@ -100,20 +102,14 @@ public class PlayerGUI implements PlayerCallbackInterface, LibraryCallbackInterf
         libraryView.removeAll();
     }
 
-    /**
-     * Populate the library with songs grouped by album.
-     *
-     * @param libraryData Map of albums with a lists of associated songs.
-     */
-    private <T> void populateLibrary(Map<T, List<Song>> libraryData) {
-        TreeMap<T, List<Song>> sortedData = new TreeMap<>(libraryData);
+    private <T extends HasSongs & Comparable<T>> void populateLibraryWithGroupedSongs(List<T> libraryData){
         DefaultTreeModel model = new DefaultTreeModel(new DefaultMutableTreeNode());
         DefaultMutableTreeNode parentNode = (DefaultMutableTreeNode) model.getRoot();
-        sortedData.forEach((k, v) -> {
-            DefaultMutableTreeNode albumNode = new DefaultMutableTreeNode(k);
-            addNodeToTreeModel(model, parentNode, albumNode);
-            Collections.sort(v);
-            v.forEach(x -> addNodeToTreeModel(model, albumNode, new DefaultMutableTreeNode(x)));
+        Collections.sort(libraryData);
+        libraryData.forEach(x -> {
+            DefaultMutableTreeNode outerNode = new DefaultMutableTreeNode(x);
+            addNodeToTreeModel(model, parentNode, outerNode);
+            x.getSongs().forEach(y -> addNodeToTreeModel(model, outerNode, new DefaultMutableTreeNode(y)));
         });
         libraryView.setModel(model);
     }
@@ -123,7 +119,7 @@ public class PlayerGUI implements PlayerCallbackInterface, LibraryCallbackInterf
      *
      * @param libraryData List of songs.
      */
-    private void populateLibrary(List<Song> libraryData) {
+    private void populateLibraryWithSongsOnly(List<Song> libraryData) {
         DefaultTreeModel model = new DefaultTreeModel(new DefaultMutableTreeNode());
         DefaultMutableTreeNode parentNode = (DefaultMutableTreeNode) model.getRoot();
         Collections.sort(libraryData);
@@ -251,9 +247,9 @@ public class PlayerGUI implements PlayerCallbackInterface, LibraryCallbackInterf
      */
     private Map<String, Runnable> createDisplayVariantMap() {
         Map<String, Runnable> value = new HashMap<>();
-        value.put("Song", () -> Gateway.listAllT(Song.class).ifPresent(this::populateLibrary));
-        value.put("Album/Song", () -> Gateway.listAllSongsGroupedByAlbum().ifPresent(this::populateLibrary));
-        value.put("Artist/Song", () -> Gateway.listAllSongsGroupedByArtist().ifPresent(this::populateLibrary));
+        value.put("Song", () -> Gateway.listAllT(Song.class).ifPresent(this::populateLibraryWithSongsOnly));
+        value.put("Album/Song", () -> Gateway.listAllT(Album.class).ifPresent(this::populateLibraryWithGroupedSongs));
+        value.put("Artist/Song", () -> Gateway.listAllT(Artist.class).ifPresent(this::populateLibraryWithGroupedSongs));
 
         return value;
     }
diff --git a/src/main/java/musicplayer/db/Gateway.java b/src/main/java/musicplayer/db/Gateway.java
index 7f3ce1f..beb8723 100644
--- a/src/main/java/musicplayer/db/Gateway.java
+++ b/src/main/java/musicplayer/db/Gateway.java
@@ -1,17 +1,14 @@
 package musicplayer.db;
 
-import musicplayer.util.AlbumArtExtractor;
 import musicplayer.model.*;
+import musicplayer.util.AlbumArtExtractor;
 import org.hibernate.Criteria;
 import org.hibernate.Session;
 import org.hibernate.criterion.Restrictions;
 
 import java.io.File;
-import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
-import java.util.stream.Collectors;
 
 /**
  * Contains database queries.
@@ -67,24 +64,6 @@ public class Gateway {
         }
     }
 
-    /**
-     * @return All songs currently in the database, grouped by album.
-     */
-    public static Optional<Map<Album, List<Song>>> listAllSongsGroupedByAlbum() {
-        Optional<List<Album>> albumList = listAllT(Album.class);
-        return (albumList.isPresent()) ? Optional.of(albumList.get().stream().collect(Collectors.toMap(x -> x, x -> new ArrayList<>(x.getSongs()))))
-                : Optional.empty();
-    }
-
-    /**
-     * @return All songs currently in the database, grouped by artist.
-     */
-    public static Optional<Map<Artist, List<Song>>> listAllSongsGroupedByArtist() {
-        Optional<List<Artist>> artistListList = listAllT(Artist.class);
-        return (artistListList.isPresent()) ? Optional.of(artistListList.get().stream().collect(Collectors.toMap(x -> x, x -> new ArrayList<>(x.getSongs()))))
-                : Optional.empty();
-    }
-
     /**
      * Add a new song to the database.
      * If the song already exists it will be updated instead.
diff --git a/src/main/java/musicplayer/model/Album.java b/src/main/java/musicplayer/model/Album.java
index a382239..7842153 100644
--- a/src/main/java/musicplayer/model/Album.java
+++ b/src/main/java/musicplayer/model/Album.java
@@ -10,7 +10,7 @@ import java.util.Optional;
 import java.util.Set;
 
 @Entity
-public class Album implements Comparable<Album>, IDBType {
+public class Album implements Comparable<Album>, IDBType, HasSongs {
 
     @Id
     @GeneratedValue(strategy = GenerationType.AUTO)
diff --git a/src/main/java/musicplayer/model/Artist.java b/src/main/java/musicplayer/model/Artist.java
index 36d5fa3..7eb7cb9 100644
--- a/src/main/java/musicplayer/model/Artist.java
+++ b/src/main/java/musicplayer/model/Artist.java
@@ -4,7 +4,7 @@ import javax.persistence.*;
 import java.util.Set;
 
 @Entity
-public class Artist implements Comparable<Artist>, IDBType {
+public class Artist implements Comparable<Artist>, IDBType, HasSongs {
 
     @Id
     @GeneratedValue(strategy = GenerationType.AUTO)
diff --git a/src/main/java/musicplayer/model/ExtractedMetadata.java b/src/main/java/musicplayer/model/ExtractedMetadata.java
index bdc1e0a..d0c73f6 100644
--- a/src/main/java/musicplayer/model/ExtractedMetadata.java
+++ b/src/main/java/musicplayer/model/ExtractedMetadata.java
@@ -3,16 +3,7 @@ package musicplayer.model;
 import org.jaudiotagger.tag.FieldKey;
 import org.jaudiotagger.tag.Tag;
 
-import javax.imageio.ImageIO;
-import java.awt.*;
-import java.awt.image.BufferedImage;
-import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
 import java.io.File;
-import java.io.IOException;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.util.Optional;
 
 /**
  * Internal class representing metadata extracted from a music file.
diff --git a/src/main/java/musicplayer/model/HasSongs.java b/src/main/java/musicplayer/model/HasSongs.java
new file mode 100644
index 0000000..836bf16
--- /dev/null
+++ b/src/main/java/musicplayer/model/HasSongs.java
@@ -0,0 +1,8 @@
+package musicplayer.model;
+
+import java.util.Set;
+
+public interface HasSongs {
+
+    Set<Song> getSongs();
+}
diff --git a/src/test/java/musicplayer/db/GatewayTest.java b/src/test/java/musicplayer/db/GatewayTest.java
index d001a25..33eb292 100644
--- a/src/test/java/musicplayer/db/GatewayTest.java
+++ b/src/test/java/musicplayer/db/GatewayTest.java
@@ -14,7 +14,6 @@ import org.junit.Test;
 
 import java.io.File;
 import java.util.List;
-import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -76,14 +75,14 @@ public class GatewayTest {
         session.save(song2);
         session.save(song3);
         session.getTransaction().commit();
-        assertTrue(Gateway.listAllSongsGroupedByArtist().isPresent());
-        Map<Artist, List<Song>> result = Gateway.listAllSongsGroupedByArtist().get();
-        assertTrue(result.size() == 2);
-        assertTrue(result.get(artist1).size() == 2);
-        assertTrue(result.get(artist2).size() == 1);
-        assertTrue(result.get(artist1).contains(song1));
-        assertTrue(result.get(artist1).contains(song2));
-        assertTrue(result.get(artist2).contains(song3));
+        List<Artist> retrievedArtists = Gateway.listAllT(Artist.class).get();
+        int index1 = retrievedArtists.indexOf(artist1);
+        int index2 = retrievedArtists.indexOf(artist2);
+        assertTrue(retrievedArtists.get(index1).getSongs().size() == 2);
+        assertTrue(retrievedArtists.get(index2).getSongs().size() == 1);
+        assertTrue(retrievedArtists.get(index1).getSongs().contains(song1));
+        assertTrue(retrievedArtists.get(index1).getSongs().contains(song2));
+        assertTrue(retrievedArtists.get(index2).getSongs().contains(song3));
     }
 
     @Test
@@ -126,19 +125,14 @@ public class GatewayTest {
         session.save(song2);
         session.save(song3);
         session.getTransaction().commit();
-        assertTrue(Gateway.listAllSongsGroupedByAlbum().isPresent());
-        Map<Album, List<Song>> result = Gateway.listAllSongsGroupedByAlbum().get();
-        assertTrue(result.size() == 2);
-        assertTrue(result.get(album1).size() == 2);
-        assertTrue(result.get(album2).size() == 1);
-        assertTrue(result.get(album1).contains(song1));
-        assertTrue(result.get(album1).contains(song2));
-        assertTrue(result.get(album2).contains(song3));
-    }
-
-    @Test
-    public void testListAllSongsGroupedByAlbumEmptyDB() throws Exception {
-        assertTrue(!Gateway.listAllSongsGroupedByAlbum().isPresent());
+        List<Album> retrievedAlbums = Gateway.listAllT(Album.class).get();
+        int index1 = retrievedAlbums.indexOf(album1);
+        int index2 = retrievedAlbums.indexOf(album2);
+        assertTrue(retrievedAlbums.get(index1).getSongs().size() == 2);
+        assertTrue(retrievedAlbums.get(index2).getSongs().size() == 1);
+        assertTrue(retrievedAlbums.get(index1).getSongs().contains(song1));
+        assertTrue(retrievedAlbums.get(index1).getSongs().contains(song2));
+        assertTrue(retrievedAlbums.get(index2).getSongs().contains(song3));
     }
 
     @Test