Minor fixes to tests

3 files changed, 12 insertions, 6 deletions

diff --git a/abcontroller_unittest.py b/abcontroller_unittest.py
index 42eaabc..ea377b7 100644
--- a/abcontroller_unittest.py
+++ b/abcontroller_unittest.py
@@ -137,3 +137,7 @@ def test_addLimitsToCurrentSongButDoNotSetCurrentLimits():
 
     checkLimits(uut, abLimits.a, abLimits.b)
     
+def test_getLimitsOfInexistentSong():
+    song = "/path/to/song"
+    uut = ABController()
+    assert uut.getLimits(song) == None
diff --git a/session_manager.py b/session_manager.py
index 057d952..52742c8 100644
--- a/session_manager.py
+++ b/session_manager.py
@@ -26,13 +26,12 @@ class SessionManager:
             
     def saveSession(self, file):
         songs = self._playlist.getSongs()
-        abLimits = self._abController.getLimits()
         session = list()
 
         for s in songs:
             entry = {
                 "path": s,
-                "ab_limits" : abLimits.get(s)
+                "ab_limits" : self._abController.getLimits(s)
             }
             session.append(entry)
 
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