diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 50ed3a3cb..3cb5e515e 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -142,7 +142,7 @@ class InfectionMonkey: return opts # TODO: By the time we finish 2292, _connect_to_island_api() may not need to return `server` - def _connect_to_island_api(self) -> Tuple[str, IIslandAPIClient]: + def _connect_to_island_api(self) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: logger.debug(f"Trying to wake up with servers: {', '.join(self._opts.servers)}") server_clients = find_available_island_apis( self._opts.servers, HTTPIslandAPIClientFactory(self._agent_event_serializer_registry) @@ -166,7 +166,7 @@ class InfectionMonkey: return server, island_api_client def _select_server( - self, server_clients: Mapping[str, IIslandAPIClient] + self, server_clients: Mapping[str, Optional[IIslandAPIClient]] ) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: for server in self._opts.servers: if server_clients[server]: