From 3bb75fd615b8d556303d69680c7dde3b9a43fde8 Mon Sep 17 00:00:00 2001
From: Till Skrodzki <till@mueskro.de>
Date: Sun, 22 Sep 2024 08:36:38 +0200
Subject: [PATCH] Renove mainscreen checks

---
 mapadroid/worker/strategy/AbstractWorkerStrategy.py    | 10 ++++++----
 .../strategy/plain/AbstractWorkerMitmStrategy.py       | 10 +++++-----
 mapadroid/worker/strategy/quest/QuestStrategy.py       |  6 +++---
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/mapadroid/worker/strategy/AbstractWorkerStrategy.py b/mapadroid/worker/strategy/AbstractWorkerStrategy.py
index 445487443..98b04bcf5 100644
--- a/mapadroid/worker/strategy/AbstractWorkerStrategy.py
+++ b/mapadroid/worker/strategy/AbstractWorkerStrategy.py
@@ -310,19 +310,21 @@ async def turn_screen_on_and_start_pogo(self):
         logger.info("turnScreenOnAndStartPogo: (Re-)Starting Pogo")
         await self.start_pogo()
 
-    async def _ensure_pogo_topmost(self):
+    async def _ensure_pogo_topmost(self, additional_eval: Optional[Callable[[], Awaitable[bool]]] = None) -> bool:
         logger.info('Checking pogo screen...')
         await self.start_pogo()
-        screen_type: ScreenType = await self._handle_screen()
+        screen_type: ScreenType = await self._handle_screen(additional_eval)
         logger.info('Checking pogo screen is finished')
-        if screen_type in [ScreenType.POGO, ScreenType.QUEST]:
+        if (screen_type in [ScreenType.POGO, ScreenType.QUEST, ScreenType.BLACK]
+                or additional_eval is not None and additional_eval()):
             return True
         else:
             return False
 
     async def _handle_screen(self, additional_eval: Optional[Callable[[], Awaitable[bool]]] = None) -> ScreenType:
         screen_type: ScreenType = ScreenType.UNDEFINED
-        while not self._worker_state.stop_worker_event.is_set() and (additional_eval is None or not await additional_eval()):
+        while not self._worker_state.stop_worker_event.is_set() and (additional_eval is None
+                                                                     or not await additional_eval()):
             if self._worker_state.login_error_count > 2:
                 logger.warning('Could not login again - clearing game data and restarting device')
                 await self.stop_pogo()
diff --git a/mapadroid/worker/strategy/plain/AbstractWorkerMitmStrategy.py b/mapadroid/worker/strategy/plain/AbstractWorkerMitmStrategy.py
index 6538cc8e7..835bbf5a3 100644
--- a/mapadroid/worker/strategy/plain/AbstractWorkerMitmStrategy.py
+++ b/mapadroid/worker/strategy/plain/AbstractWorkerMitmStrategy.py
@@ -2,11 +2,12 @@
 import math
 import time
 from abc import ABC, abstractmethod
-from typing import Optional, Set, Tuple, Union, Any
+from typing import Any, Optional, Set, Tuple, Union
 
 from google.protobuf.internal.containers import RepeatedCompositeFieldContainer
 from redis import Redis
 
+import mapadroid.mitm_receiver.protos.Rpc_pb2 as pogoprotos
 from mapadroid.data_handler.mitm_data.holder.latest_mitm_data.LatestMitmDataEntry import \
     LatestMitmDataEntry
 from mapadroid.mapping_manager.MappingManagerDevicemappingKey import \
@@ -22,7 +23,6 @@
 from mapadroid.worker.ReceivedTypeEnum import ReceivedType
 from mapadroid.worker.strategy.AbstractMitmBaseStrategy import \
     AbstractMitmBaseStrategy
-import mapadroid.mitm_receiver.protos.Rpc_pb2 as pogoprotos
 
 logger = get_logger(LoggerEnums.worker)
 
@@ -42,9 +42,9 @@ async def pre_work_loop(self):
         if not await self._wait_for_injection() or self._worker_state.stop_worker_event.is_set():
             raise InternalStopWorkerException("Worker stopped in pre work loop")
 
-        reached_main_menu = await self._check_pogo_main_screen(10, True)
-        if not reached_main_menu:
-            logger.info("Main menu was not reached, trying to restart pogo")
+        pogo_topmost = await self._communicator.is_pogo_topmost()
+        if not pogo_topmost:
+            logger.info("Pogo is not topmost app, trying to restart pogo")
             if not await self._restart_pogo():
                 # TODO: put in loop, count up for a reboot ;)
                 raise InternalStopWorkerException("Worker stopped in pre work loop")
diff --git a/mapadroid/worker/strategy/quest/QuestStrategy.py b/mapadroid/worker/strategy/quest/QuestStrategy.py
index b54f89acd..af4168ebd 100644
--- a/mapadroid/worker/strategy/quest/QuestStrategy.py
+++ b/mapadroid/worker/strategy/quest/QuestStrategy.py
@@ -496,11 +496,11 @@ async def switch_account(self):
             raise InternalStopWorkerException("Failed switching accounts")
         else:
             await asyncio.sleep(10)
-            reached_main_menu = await self._check_pogo_main_screen(10, True)
-            if not reached_main_menu:
+            pogo_topmost = await self._ensure_pogo_topmost(additional_eval=self._is_injected)
+            if not pogo_topmost:
                 if not await self._restart_pogo():
                     # TODO: put in loop, count up for a reboot ;)
-                    raise InternalStopWorkerException("Failed reaching the pogo main screen after switching accounts")
+                    raise InternalStopWorkerException("Failed having pogo as topmost app after switching accounts")
 
     async def _get_ids_iv_and_scanmode(self) -> Tuple[List[int], str]:
         injected_settings = {}