learning-loop-node 0.10.13__py3-none-any.whl → 0.10.14__py3-none-any.whl

learning_loop_node/detector/rest/outbox_mode.py

@@ -6,10 +6,13 @@ from ..outbox import Outbox
 router = APIRouter()
 
 
-@router.get("/outbox_mode")
+@router.get("/outbox_mode", response_class=PlainTextResponse)
 async def get_outbox_mode(request: Request):
     '''
+    Get the outbox mode of the detector node.
+
     Example Usage
+
         curl http://localhost/outbox_mode
     '''
     outbox: Outbox = request.app.outbox
@@ -19,7 +22,10 @@ async def get_outbox_mode(request: Request):
 @router.put("/outbox_mode")
 async def put_outbox_mode(request: Request):
     '''
+    Set the outbox mode of the detector node.
+
     Example Usage
+
         curl -X PUT -d "continuous_upload" http://localhost/outbox_mode
         curl -X PUT -d "stopped" http://localhost/outbox_mode
     '''

learning_loop_node/node.py

@@ -8,7 +8,6 @@ from datetime import datetime
 from typing import Any, Optional
 
 import aiohttp
-import socketio
 from aiohttp import TCPConnector
 from fastapi import FastAPI
 from socketio import AsyncClient
@@ -21,6 +20,10 @@ from .loop_communication import LoopCommunicator
 from .rest import router
 
 
+class NodeConnectionError(Exception):
+    pass
+
+
 class Node(FastAPI):
 
     def __init__(self, name: str, uuid: Optional[str] = None, node_type: str = 'node', needs_login: bool = True):
@@ -43,8 +46,7 @@ class Node(FastAPI):
         self.needs_login = needs_login
 
         self.log = logging.getLogger('Node')
-        self.loop_communicator = LoopCommunicator()
-        self.websocket_url = self.loop_communicator.websocket_url()
+        self.init_loop_communicator()
         self.data_exchanger = DataExchanger(None, self.loop_communicator)
 
         self.startup_datetime = datetime.now()
@@ -56,9 +58,20 @@ class Node(FastAPI):
                             'nodeType': node_type}
 
         self.repeat_task: Any = None
+        self.socket_connection_broken = False
+        self._skip_repeat_loop = False
 
         self.include_router(router)
 
+        self.CONNECTED_TO_LOOP = asyncio.Event()
+        self.DISCONNECTED_FROM_LOOP = asyncio.Event()
+
+        self.repeat_loop_lock = asyncio.Lock()
+
+    def init_loop_communicator(self):
+        self.loop_communicator = LoopCommunicator()
+        self.websocket_url = self.loop_communicator.websocket_url()
+
     @property
     def sio_client(self) -> AsyncClient:
         if self._sio_client is None:
@@ -69,7 +82,10 @@ class Node(FastAPI):
     @asynccontextmanager
     async def lifespan(self, app: FastAPI):  # pylint: disable=unused-argument
         try:
-            await self._on_startup()
+            try:
+                await self._on_startup()
+            except Exception:
+                self.log.exception('Fatal error during startup: %s')
             self.repeat_task = asyncio.create_task(self.repeat_loop())
             yield
         finally:
@@ -83,13 +99,10 @@ class Node(FastAPI):
 
     async def _on_startup(self):
         self.log.info('received "startup" lifecycle-event')
-        # activate_asyncio_warnings()
-        if self.needs_login:
-            await self.loop_communicator.backend_ready()
-            self.log.info('ensuring login')
-            await self.loop_communicator.ensure_login()
-        self.log.info('create sio client')
-        await self.create_sio_client()
+        try:
+            await self.reconnect_to_loop()
+        except Exception:
+            self.log.warning('Could not establish sio connection to loop during startup')
         self.log.info('done')
         await self.on_startup()
 
@@ -102,55 +115,93 @@ class Node(FastAPI):
         await self.on_shutdown()
 
     async def repeat_loop(self) -> None:
-        """NOTE: with the lifespan approach, we cannot use @repeat_every anymore :("""
         while True:
+            if self._skip_repeat_loop:
+                self.log.debug('node is muted, skipping repeat loop')
+                await asyncio.sleep(1)
+                continue
             try:
-                await self._on_repeat()
+                async with self.repeat_loop_lock:
+                    await self._ensure_sio_connection()
+                    await self.on_repeat()
             except asyncio.CancelledError:
                 return
-            except Exception as e:
-                self.log.exception(f'error in repeat loop: {e}')
+            except Exception:
+                self.log.exception('error in repeat loop')
+
             await asyncio.sleep(5)
 
-    async def _on_repeat(self):
-        if not self.sio_client.connected:
-            self.log.info('Reconnecting to loop via sio')
-            await self.connect_sio()
-            if not self.sio_client.connected:
-                self.log.warning('Could not connect to loop via sio')
-                return
-        await self.on_repeat()
+    async def _ensure_sio_connection(self):
+        if self.socket_connection_broken or self._sio_client is None or not self.sio_client.connected:
+            self.log.info('Reconnecting to loop via sio due to %s',
+                          'broken connection' if self.socket_connection_broken else 'no connection')
+            await self.reconnect_to_loop()
+
+    async def reconnect_to_loop(self):
+        """Initialize the loop communicator, log in if needed and reconnect to the loop via socket.io."""
+        self.init_loop_communicator()
+        if self.needs_login:
+            await self.loop_communicator.ensure_login(relogin=True)
+        try:
+            await self._reconnect_socketio()
+        except NodeConnectionError:
+            self.log.exception('Could not reset sio connection to loop')
+            self.socket_connection_broken = True
+            raise
+
+        self.socket_connection_broken = False
+
+    def set_skip_repeat_loop(self, value: bool):
+        self._skip_repeat_loop = value
+        self.log.info('node is muted: %s', value)
 
     # --------------------------------------------------- SOCKET.IO ---------------------------------------------------
 
-    async def create_sio_client(self):
-        """Create a socket.io client that communicates with the learning loop and register the events.
-        Note: The method is called in startup and soft restart of detector, so the _sio_client should always be available."""
+    async def _reconnect_socketio(self):
+        """Create a socket.io client, connect it to the learning loop and register its events.
+        The current client is disconnected and deleted if it already exists."""
+
+        self.log.debug('-------------- Connecting to loop via socket.io -------------------')
+        self.log.debug('HTTP Cookies: %s\n', self.loop_communicator.get_cookies())
 
+        if self._sio_client is not None:
+            try:
+                await self.sio_client.disconnect()
+                self.log.info('disconnected from loop via sio')
+                # NOTE: without waiting for the disconnect event, we might disconnect the next connection too early
+                await asyncio.wait_for(self.DISCONNECTED_FROM_LOOP.wait(), timeout=5)
+            except asyncio.TimeoutError:
+                self.log.warning(
+                    'Did not receive disconnect event from loop within 5 seconds.\nContinuing with new connection...')
+            except Exception as e:
+                self.log.warning('Could not disconnect from loop via sio: %s.\nIgnoring...', e)
+            self._sio_client = None
+
+        connector = None
         if self.loop_communicator.ssl_cert_path:
-            logging.info(f'SIO using SSL certificate path: {self.loop_communicator.ssl_cert_path}')
+            logging.info('SIO using SSL certificate path: %s', self.loop_communicator.ssl_cert_path)
             ssl_context = ssl.create_default_context(cafile=self.loop_communicator.ssl_cert_path)
             ssl_context.check_hostname = False
             ssl_context.verify_mode = ssl.CERT_REQUIRED
             connector = TCPConnector(ssl=ssl_context)
-            self._sio_client = AsyncClient(request_timeout=20,
-                                           http_session=aiohttp.ClientSession(cookies=self.loop_communicator.get_cookies(),
-                                                                              connector=connector))
 
-        else:
-            self._sio_client = AsyncClient(request_timeout=20,
-                                           http_session=aiohttp.ClientSession(cookies=self.loop_communicator.get_cookies()))
+        self._sio_client = AsyncClient(request_timeout=20, http_session=aiohttp.ClientSession(
+            cookies=self.loop_communicator.get_cookies(), connector=connector))
 
         # pylint: disable=protected-access
-        self.sio_client._trigger_event = ensure_socket_response(self.sio_client._trigger_event)
+        self._sio_client._trigger_event = ensure_socket_response(self._sio_client._trigger_event)
 
         @self._sio_client.event
         async def connect():
             self.log.info('received "connect" via sio from loop.')
+            self.CONNECTED_TO_LOOP.set()
+            self.DISCONNECTED_FROM_LOOP.clear()
 
         @self._sio_client.event
         async def disconnect():
             self.log.info('received "disconnect" via sio from loop.')
+            self.DISCONNECTED_FROM_LOOP.set()
+            self.CONNECTED_TO_LOOP.clear()
 
         @self._sio_client.event
         async def restart():
@@ -158,21 +209,15 @@ class Node(FastAPI):
             sys.exit(0)
 
         self.register_sio_events(self._sio_client)
-
-    async def connect_sio(self):
         try:
-            await self.sio_client.disconnect()
-        except Exception:
-            pass
-
-        self.log.info(f'(re)connecting to Learning Loop at {self.websocket_url}')
-        try:
-            await self.sio_client.connect(f"{self.websocket_url}", headers=self.sio_headers, socketio_path="/ws/socket.io")
-            self.log.info('connected to Learning Loop')
-        except socketio.exceptions.ConnectionError:  # type: ignore
-            self.log.warning('connection error')
-        except Exception:
-            self.log.exception(f'error while connecting to "{self.websocket_url}". Exception:')
+            await self._sio_client.connect(f"{self.websocket_url}", headers=self.sio_headers, socketio_path="/ws/socket.io")
+        except Exception as e:
+            self.log.exception('Could not connect socketio client to loop')
+            raise NodeConnectionError('Could not connect socketio client to loop') from e
+
+        if not self._sio_client.connected:
+            self.log.exception('Could not connect socketio client to loop')
+            raise NodeConnectionError('Could not connect socketio client to loop')
 
     # --------------------------------------------------- ABSTRACT METHODS ---------------------------------------------------
 

learning_loop_node/rest.py

@@ -1,7 +1,7 @@
 import logging
 from typing import TYPE_CHECKING
 
-from fastapi import APIRouter, Request, HTTPException
+from fastapi import APIRouter, HTTPException, Request
 
 if TYPE_CHECKING:
     from .node import Node
@@ -12,7 +12,7 @@ logger = logging.getLogger('Node.rest')
 
 
 @router.put("/debug_logging")
-async def _debug_logging(request: Request):
+async def _debug_logging(request: Request) -> str:
     '''
     Example Usage
 
@@ -30,3 +30,26 @@ async def _debug_logging(request: Request):
         node.log.setLevel('DEBUG')
         return 'on'
     raise HTTPException(status_code=400, detail='Invalid state')
+
+
+@router.put("/socketio")
+async def _socketio(request: Request) -> str:
+    '''
+    Enable or disable the socketio connection to the learning loop.
+    Not intended to be used outside of testing.
+
+    Example Usage
+
+        curl -X PUT -d "on" http://hosturl/socketio
+    '''
+    state = str(await request.body(), 'utf-8')
+    node: 'Node' = request.app
+
+    if state == 'off':
+        await node.sio_client.disconnect()
+        node.set_skip_repeat_loop(True)  # Prevent auto-reconnection
+        return 'off'
+    if state == 'on':
+        node.set_skip_repeat_loop(False)  # Allow auto-reconnection (1 sec delay)
+        return 'on'
+    raise HTTPException(status_code=400, detail='Invalid state')

learning_loop_node/tests/detector/conftest.py

@@ -41,8 +41,8 @@ def should_have_segmentations(request) -> bool:
 async def test_detector_node():
     """Initializes and runs a detector testnode. Note that the running instance and the one the function returns are not the same instances!"""
 
-    os.environ['ORGANIZATION'] = 'zauberzeug'
-    os.environ['PROJECT'] = 'demo'
+    os.environ['LOOP_ORGANIZATION'] = 'zauberzeug'
+    os.environ['LOOP_PROJECT'] = 'demo'
 
     detector = TestingDetectorLogic()
     node = DetectorNode(name='test', detector=detector)
@@ -143,8 +143,8 @@ def mock_detector_logic():
 
 @pytest.fixture
 def detector_node(mock_detector_logic):
-    os.environ['ORGANIZATION'] = 'test_organization'
-    os.environ['PROJECT'] = 'test_project'
+    os.environ['LOOP_ORGANIZATION'] = 'test_organization'
+    os.environ['LOOP_PROJECT'] = 'test_project'
     return DetectorNode(name="test_node", detector=mock_detector_logic)
 
 # ====================================== REDUNDANT FIXTURES IN ALL CONFTESTS ! ======================================

learning_loop_node/tests/detector/test_client_communication.py

@@ -93,10 +93,10 @@ async def test_sio_upload(test_detector_node: DetectorNode, sio_client):
 
 # NOTE: This test seems to be flaky.
 async def test_about_endpoint(test_detector_node: DetectorNode):
-    await asyncio.sleep(3)
+    await asyncio.sleep(11)
     response = requests.get(f'http://localhost:{GLOBALS.detector_port}/about', timeout=30)
 
-    assert response.status_code == 200
+    assert response.status_code == 200, response.content
     response_dict = json.loads(response.content)
     assert response_dict['model_info']
     model_information = ModelInformation.from_dict(response_dict['model_info'])
@@ -108,59 +108,60 @@ async def test_about_endpoint(test_detector_node: DetectorNode):
 
 
 async def test_model_version_api(test_detector_node: DetectorNode):
-    await asyncio.sleep(3)
+    await asyncio.sleep(11)
 
     response = requests.get(f'http://localhost:{GLOBALS.detector_port}/model_version', timeout=30)
-    assert response.status_code == 200
+    assert response.status_code == 200, response.content
     response_dict = json.loads(response.content)
+    assert response_dict['version_control'] == 'follow_loop'
     assert response_dict['current_version'] == '1.1'
     assert response_dict['target_version'] == '1.1'
     assert response_dict['loop_version'] == '1.1'
     assert response_dict['local_versions'] == ['1.1']
-    assert response_dict['version_control'] == 'follow_loop'
 
     response = requests.put(f'http://localhost:{GLOBALS.detector_port}/model_version', data='1.0', timeout=30)
+    assert response.status_code == 200, response.content
     response = requests.get(f'http://localhost:{GLOBALS.detector_port}/model_version', timeout=30)
-    assert response.status_code == 200
+    assert response.status_code == 200, response.content
     response_dict = json.loads(response.content)
+    assert response_dict['version_control'] == 'specific_version'
     assert response_dict['current_version'] == '1.1'
     assert response_dict['target_version'] == '1.0'
     assert response_dict['loop_version'] == '1.1'
     assert response_dict['local_versions'] == ['1.1']
-    assert response_dict['version_control'] == 'specific_version'
 
     await asyncio.sleep(11)
-
     response = requests.get(f'http://localhost:{GLOBALS.detector_port}/model_version', timeout=30)
-    assert response.status_code == 200
+    assert response.status_code == 200, response.content
     response_dict = json.loads(response.content)
+    assert response_dict['version_control'] == 'specific_version'
     assert response_dict['current_version'] == '1.0'
     assert response_dict['target_version'] == '1.0'
     assert response_dict['loop_version'] == '1.1'
     assert set(response_dict['local_versions']) == set(['1.1', '1.0'])
-    assert response_dict['version_control'] == 'specific_version'
 
     response = requests.put(f'http://localhost:{GLOBALS.detector_port}/model_version', data='pause', timeout=30)
+    assert response.status_code == 200, response.content
     await asyncio.sleep(11)
     response = requests.get(f'http://localhost:{GLOBALS.detector_port}/model_version', timeout=30)
-    assert response.status_code == 200
+    assert response.status_code == 200, response.content
     response_dict = json.loads(response.content)
+    assert response_dict['version_control'] == 'pause'
     assert response_dict['current_version'] == '1.0'
     assert response_dict['target_version'] == '1.0'
     assert response_dict['loop_version'] == '1.1'
     assert set(response_dict['local_versions']) == set(['1.1', '1.0'])
-    assert response_dict['version_control'] == 'pause'
 
     response = requests.put(f'http://localhost:{GLOBALS.detector_port}/model_version', data='follow_loop', timeout=30)
     await asyncio.sleep(11)
     response = requests.get(f'http://localhost:{GLOBALS.detector_port}/model_version', timeout=30)
-    assert response.status_code == 200
+    assert response.status_code == 200, response.content
     response_dict = json.loads(response.content)
+    assert response_dict['version_control'] == 'follow_loop'
     assert response_dict['current_version'] == '1.1'
     assert response_dict['target_version'] == '1.1'
     assert response_dict['loop_version'] == '1.1'
     assert set(response_dict['local_versions']) == set(['1.1', '1.0'])
-    assert response_dict['version_control'] == 'follow_loop'
 
 
 async def test_rest_outbox_mode(test_detector_node: DetectorNode):
@@ -169,9 +170,9 @@ async def test_rest_outbox_mode(test_detector_node: DetectorNode):
     def check_switch_to_mode(mode: str):
         response = requests.put(f'http://localhost:{GLOBALS.detector_port}/outbox_mode',
                                 data=mode, timeout=30)
-        assert response.status_code == 200
+        assert response.status_code == 200, response.content
         response = requests.get(f'http://localhost:{GLOBALS.detector_port}/outbox_mode', timeout=30)
-        assert response.status_code == 200
+        assert response.status_code == 200, response.content
         assert response.content == mode.encode()
 
     check_switch_to_mode('stopped')
@@ -185,7 +186,7 @@ async def test_api_responsive_during_large_upload(test_detector_node: DetectorNo
     with open(test_image_path, 'rb') as f:
         image_bytes = f.read()
 
-    for _ in range(100):
+    for _ in range(200):
         test_detector_node.outbox.save(image_bytes)
 
     outbox_size_early = len(get_outbox_files(test_detector_node.outbox))
@@ -193,8 +194,9 @@ async def test_api_responsive_during_large_upload(test_detector_node: DetectorNo
 
     # check if api is still responsive
     response = requests.get(f'http://localhost:{GLOBALS.detector_port}/outbox_mode', timeout=2)
-    assert response.status_code == 200
+    assert response.status_code == 200, response.content
 
     await asyncio.sleep(5)
     outbox_size_late = len(get_outbox_files(test_detector_node.outbox))
-    assert outbox_size_early > outbox_size_late > 0, 'The outbox should have been partially emptied.'
+    assert outbox_size_late > 0, 'The outbox should not be fully cleared, maybe the node was too fast.'
+    assert outbox_size_early > outbox_size_late, 'The outbox should have been partially emptied.'

learning_loop_node/tests/detector/test_detector_node.py

@@ -57,9 +57,9 @@ async def test_get_detections(detector_node: DetectorNode, monkeypatch):
 
         # Check if detections were processed
         assert result is not None
-        assert "box_detections" in result
-        assert len(result["box_detections"]) == 1
-        assert result["box_detections"][0]["category_name"] == "test"
+        assert result.box_detections is not None
+        assert len(result.box_detections) == 1
+        assert result.box_detections[0].category_name == "test"
 
         # Check if the correct upload method was called
         assert filtered_upload_called == expect_filtered

learning_loop_node/tests/trainer/conftest.py

@@ -1,7 +1,6 @@
-from ...globals import GLOBALS
-import shutil
 import logging
 import os
+import shutil
 import socket
 from multiprocessing import log_to_stderr
 
@@ -9,6 +8,7 @@ import icecream
 import pytest
 
 from ...data_classes import Context
+from ...globals import GLOBALS
 from ...trainer.trainer_node import TrainerNode
 from .testing_trainer_logic import TestingTrainerLogic
 
@@ -23,8 +23,8 @@ icecream.install()
 
 @pytest.fixture()
 async def test_initialized_trainer_node():
-    os.environ['ORGANIZATION'] = 'zauberzeug'
-    os.environ['PROJECT'] = 'demo'
+    os.environ['LOOP_ORGANIZATION'] = 'zauberzeug'
+    os.environ['LOOP_PROJECT'] = 'demo'
 
     trainer = TestingTrainerLogic()
     node = TrainerNode(name='test', trainer_logic=trainer, uuid='NOD30000-0000-0000-0000-000000000000')

learning_loop_node/tests/trainer/states/test_state_detecting.py

@@ -7,11 +7,10 @@ from ..state_helper import assert_training_state, create_active_training_file
 from ..testing_trainer_logic import TestingTrainerLogic
 
 # pylint: disable=protected-access
-error_key = 'detecting'
 
 
-def trainer_has_error(trainer: TrainerLogic):
-    return trainer.errors.has_error_for(error_key)
+def trainer_has_detecting_error(trainer: TrainerLogic):
+    return trainer.errors.has_error_for('detecting')
 
 
 async def test_successful_detecting(test_initialized_trainer: TestingTrainerLogic):
@@ -25,7 +24,7 @@ async def test_successful_detecting(test_initialized_trainer: TestingTrainerLogi
     await assert_training_state(trainer.training, TrainerState.Detecting, timeout=1, interval=0.001)
     await assert_training_state(trainer.training, TrainerState.Detected, timeout=10, interval=0.001)
 
-    assert trainer_has_error(trainer) is False
+    assert trainer_has_detecting_error(trainer) is False
     assert trainer.training.training_state == TrainerState.Detected
     assert trainer.node.last_training_io.load() == trainer.training
     assert trainer.active_training_io.detections_exist()
@@ -37,7 +36,7 @@ async def test_detecting_can_be_aborted(test_initialized_trainer: TestingTrainer
     trainer._init_from_last_training()
     trainer.training.model_uuid_for_detecting = '12345678-bobo-7e92-f95f-424242424242'
 
-    _ = asyncio.get_running_loop().create_task(trainer._run())
+    trainer._begin_training_task()
 
     await assert_training_state(trainer.training, TrainerState.Detecting, timeout=5, interval=0.001)
     await trainer.stop()
@@ -54,13 +53,13 @@ async def test_model_not_downloadable_error(test_initialized_trainer: TestingTra
                                 model_uuid_for_detecting='00000000-0000-0000-0000-000000000000')  # bad model id
     trainer._init_from_last_training()
 
-    _ = asyncio.get_running_loop().create_task(trainer._run())
+    trainer._begin_training_task()
 
-    await assert_training_state(trainer.training, 'detecting', timeout=1, interval=0.001)
-    await assert_training_state(trainer.training, 'train_model_uploaded', timeout=1, interval=0.001)
+    await assert_training_state(trainer.training, TrainerState.Detecting, timeout=1, interval=0.001)
+    await assert_training_state(trainer.training, TrainerState.TrainModelUploaded, timeout=5, interval=0.001)
     await asyncio.sleep(0.1)
 
-    assert trainer_has_error(trainer)
+    assert trainer_has_detecting_error(trainer)
     assert trainer.training.training_state == TrainerState.TrainModelUploaded
     assert trainer.training.model_uuid_for_detecting == '00000000-0000-0000-0000-000000000000'
     assert trainer.node.last_training_io.load() == trainer.training

learning_loop_node/tests/trainer/states/test_state_download_train_model.py

@@ -20,8 +20,8 @@ async def test_downloading_is_successful(test_initialized_trainer: TestingTraine
         trainer._perform_state('download_model',
                                TrainerState.TrainModelDownloading,
                                TrainerState.TrainModelDownloaded, trainer._download_model))
-    await assert_training_state(trainer.training, 'train_model_downloading', timeout=1, interval=0.001)
-    await assert_training_state(trainer.training, 'train_model_downloaded', timeout=1, interval=0.001)
+    await assert_training_state(trainer.training, TrainerState.TrainModelDownloading, timeout=1, interval=0.001)
+    await assert_training_state(trainer.training, TrainerState.TrainModelDownloaded, timeout=10, interval=0.001)
 
     assert trainer.training.training_state == TrainerState.TrainModelDownloaded
     assert trainer.node.last_training_io.load() == trainer.training
@@ -34,11 +34,11 @@ async def test_downloading_is_successful(test_initialized_trainer: TestingTraine
 
 async def test_abort_download_model(test_initialized_trainer: TestingTrainerLogic):
     trainer = test_initialized_trainer
-    create_active_training_file(trainer, training_state='data_downloaded')
+    create_active_training_file(trainer, training_state=TrainerState.DataDownloaded)
     trainer._init_from_last_training()
 
-    _ = asyncio.get_running_loop().create_task(trainer._run())
-    await assert_training_state(trainer.training, 'train_model_downloading', timeout=1, interval=0.001)
+    trainer._begin_training_task()
+    await assert_training_state(trainer.training, TrainerState.TrainModelDownloading, timeout=1, interval=0.001)
 
     await trainer.stop()
     await asyncio.sleep(0.1)
@@ -53,9 +53,9 @@ async def test_downloading_failed(test_initialized_trainer: TestingTrainerLogic)
                                 base_model_uuid_or_name='00000000-0000-0000-0000-000000000000')  # bad model id)
     trainer._init_from_last_training()
 
-    _ = asyncio.get_running_loop().create_task(trainer._run())
-    await assert_training_state(trainer.training, 'train_model_downloading', timeout=1, interval=0.001)
-    await assert_training_state(trainer.training, TrainerState.DataDownloaded, timeout=1, interval=0.001)
+    trainer._begin_training_task()
+    await assert_training_state(trainer.training, TrainerState.TrainModelDownloading, timeout=1, interval=0.001)
+    await assert_training_state(trainer.training, TrainerState.DataDownloaded, timeout=10, interval=0.001)
 
     assert trainer.errors.has_error_for('download_model')
     assert trainer._training is not None  # pylint: disable=protected-access

learning_loop_node/tests/trainer/states/test_state_prepare.py

@@ -6,11 +6,10 @@ from ..state_helper import assert_training_state, create_active_training_file
 from ..testing_trainer_logic import TestingTrainerLogic
 
 # pylint: disable=protected-access
-error_key = 'prepare'
 
 
-def trainer_has_error(trainer: TrainerLogic):
-    return trainer.errors.has_error_for(error_key)
+def trainer_has_prepare_error(trainer: TrainerLogic):
+    return trainer.errors.has_error_for('prepare')
 
 
 async def test_preparing_is_successful(test_initialized_trainer: TestingTrainerLogic):
@@ -19,7 +18,7 @@ async def test_preparing_is_successful(test_initialized_trainer: TestingTrainerL
     trainer._init_from_last_training()
 
     await trainer._perform_state('prepare', TrainerState.DataDownloading, TrainerState.DataDownloaded, trainer._prepare)
-    assert trainer_has_error(trainer) is False
+    assert trainer_has_prepare_error(trainer) is False
     assert trainer.training.training_state == TrainerState.DataDownloaded
     assert trainer.training.data is not None
     assert trainer.node.last_training_io.load() == trainer.training
@@ -30,7 +29,7 @@ async def test_abort_preparing(test_initialized_trainer: TestingTrainerLogic):
     create_active_training_file(trainer)
     trainer._init_from_last_training()
 
-    _ = asyncio.get_running_loop().create_task(trainer._run())
+    trainer._begin_training_task()
     await assert_training_state(trainer.training, TrainerState.DataDownloading, timeout=1, interval=0.001)
 
     await trainer.stop()
@@ -48,9 +47,9 @@ async def test_request_error(test_initialized_trainer: TestingTrainerLogic):
 
     _ = asyncio.get_running_loop().create_task(trainer._run())
     await assert_training_state(trainer.training, TrainerState.DataDownloading, timeout=3, interval=0.001)
-    await assert_training_state(trainer.training, TrainerState.Initialized, timeout=3, interval=0.001)
+    await assert_training_state(trainer.training, TrainerState.Initialized, timeout=10, interval=0.001)
 
-    assert trainer_has_error(trainer)
+    assert trainer_has_prepare_error(trainer)
     assert trainer._training is not None  # pylint: disable=protected-access
     assert trainer.training.training_state == TrainerState.Initialized
     assert trainer.node.last_training_io.load() == trainer.training