Minor fixes to tests

1 files changed, 7 insertions, 4 deletions

diff --git a/session_manager_unittest.py b/session_manager_unittest.py
index 0169132..7fb8c55 100644
--- a/session_manager_unittest.py
+++ b/session_manager_unittest.py
@@ -42,8 +42,8 @@ class ABControllerMock:
             self.limits[song] = list()
         self.limits[song].append([aLimit, bLimit])
     
-    def getLimits(self):
-        return self.limits
+    def getLimits(self, song):
+        return self.limits.get(song)
 
 class MockFile:
     def __init__(self, init=""):
@@ -133,5 +133,8 @@ def test_loadAndSaveEmptySession():
     assert loads(sessionFile.read()) == list()
 
     uut.loadSession(sessionFile)
-    assert playlistMock.getSongs() == list()
-    assert abControllerMock.getLimits() == dict()
+
+    songs = playlistMock.getSongs()
+    assert songs == list()
+    for s in songs:
+        assert abControllerMock.getLimits(s) == None