diff --git a/spectur_live_view.py b/spectur_live_view.py
index e81812a..caab3ea 100644
--- a/spectur_live_view.py
+++ b/spectur_live_view.py
@@ -4,7 +4,6 @@ Open an automated selenium webdriver, and download images from Live View mode.
 """
 
 import os
-import sys
 import glob
 import time
 import pytz
@@ -175,7 +174,7 @@ def get_next_image(driver, cookies):
 
     # Check if recording has completed
     if t[-1] - t[0] > RECORDING_DURATION:
-        sys.exit()
+        raise(StopIteration)
 
     # Count images collected
     n = len(t)
@@ -211,6 +210,9 @@ def main():
             # Attempt to re-login if server goes down
             time.sleep(5)
             cookies = login(driver)
+        except StopIteration:
+            # Stop when recording has completed
+            break
 
 
 # Create counter