From 7beb10df8d06b7974b39f3b77e3aeaccd307fd59 Mon Sep 17 00:00:00 2001
From: DevilXD <bartosz.peszek@gmail.com>
Date: Mon, 21 Mar 2022 13:22:22 +0100
Subject: [PATCH] Fix not being able to close the app during login flow

---
 exceptions.py | 25 +++++++++++++++++++++++++
 gui.py        |  9 ++++++++-
 twitch.py     |  8 ++++++--
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/exceptions.py b/exceptions.py
index ab20a3d..08b9927 100644
--- a/exceptions.py
+++ b/exceptions.py
@@ -1,4 +1,7 @@
 class MinerException(Exception):
+    """
+    Base exception class for this application.
+    """
     def __init__(self, *args: object):
         if args:
             super().__init__(*args)
@@ -6,7 +9,23 @@ class MinerException(Exception):
             super().__init__("Unknown miner error")
 
 
+class ExitRequest(MinerException):
+    """
+    Raised when the application is requested to exit from outside of the main loop.
+
+    Intended for internal use only.
+    """
+    def __init__(self, *args: object):
+        if args:
+            super().__init__(*args)
+        else:
+            super().__init__("Unknown error during request")
+
+
 class RequestException(MinerException):
+    """
+    Raised for cases where a web request doesn't return what we wanted it to.
+    """
     def __init__(self, *args: object):
         if args:
             super().__init__(*args)
@@ -15,6 +34,9 @@ class RequestException(MinerException):
 
 
 class LoginException(RequestException):
+    """
+    Raised when an exception occurs during login phase.
+    """
     def __init__(self, *args: object):
         if args:
             super().__init__(*args)
@@ -23,5 +45,8 @@ class LoginException(RequestException):
 
 
 class CaptchaRequired(LoginException):
+    """
+    The most dreaded thing about automated scripts.
+    """
     def __init__(self):
         super().__init__("Captcha is required")
diff --git a/gui.py b/gui.py
index 2be628c..1e8b535 100644
--- a/gui.py
+++ b/gui.py
@@ -16,6 +16,7 @@ except ModuleNotFoundError as exc:
     raise ImportError("You have to run 'pip install pystray' first") from exc
 
 from utils import resource_path
+from exceptions import ExitRequest
 from registry import RegistryKey, ValueType
 from constants import SELF_PATH, FORMATTER, WS_TOPICS_LIMIT, MAX_WEBSOCKETS, WINDOW_TITLE, State
 
@@ -333,7 +334,13 @@ class LoginForm:
         self._manager.print("Please log in to continue.")
         self._confirm.clear()
         self._button.config(state="normal")
-        await self._confirm.wait()
+        # NOTE: we need this to allow for the closing window event to break the waiting here
+        await asyncio.wait(
+            [self._confirm.wait(), self._manager.wait_until_closed()],
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+        if self._manager.close_requested:
+            raise ExitRequest()
         self._button.config(state="disabled")
         data = LoginData(self._login_entry.get(), self._pass_entry.get(), self._token_entry.get())
         return data
diff --git a/twitch.py b/twitch.py
index 36c24ee..f1a10de 100644
--- a/twitch.py
+++ b/twitch.py
@@ -20,7 +20,7 @@ from channel import Channel
 from websocket import WebsocketPool
 from inventory import DropsCampaign
 from utils import task_wrapper, timestamp, AwaitableValue, OrderedSet
-from exceptions import RequestException, LoginException, CaptchaRequired
+from exceptions import ExitRequest, RequestException, LoginException, CaptchaRequired
 from constants import (
     BASE_URL,
     CLIENT_ID,
@@ -162,7 +162,11 @@ class Twitch:
         • Changing the stream that's being watched if necessary
         """
         self.gui.start()
-        await self.check_login()
+        try:
+            await self.check_login()
+        except ExitRequest:
+            # we've been requested to exit during login, most likely
+            return
         # NOTE: watch task is explicitly restarted on each new run
         if self._watching_task is not None:
             self._watching_task.cancel()