learning-loop-node 0.9.2__py3-none-any.whl → 0.10.0__py3-none-any.whl

learning_loop_node/helpers/misc.py

@@ -1,14 +1,21 @@
 """original copied from https://quantlane.com/blog/ensure-asyncio-task-exceptions-get-logged/"""
 import asyncio
 import functools
+import json
 import logging
 import os
+import shutil
+import sys
 from dataclasses import asdict
+from glob import glob
+from time import perf_counter
 from typing import Any, Coroutine, List, Optional, Tuple, TypeVar
+from uuid import UUID, uuid4
 
 import pynvml
 
-from ..data_classes import SocketResponse
+from ..data_classes import Context, SocketResponse, Training
+from ..globals import GLOBALS
 
 T = TypeVar('T')
 
@@ -48,7 +55,7 @@ def _handle_task_result(task: asyncio.Task,
         logger.exception(message, *message_args)
 
 
-def get_free_memory_mb() -> float:  # TODO check if this is used
+def get_free_memory_mb() -> float:  # NOTE used by yolov5
     pynvml.nvmlInit()
     h = pynvml.nvmlDeviceGetHandleByIndex(0)
     info = pynvml.nvmlDeviceGetMemoryInfo(h)
@@ -56,16 +63,33 @@ def get_free_memory_mb() -> float:  # TODO check if this is used
     return free
 
 
+async def is_valid_image(filename: str, check_jpeg: bool) -> bool:
+    if not os.path.isfile(filename) or os.path.getsize(filename) == 0:
+        return False
+    if not check_jpeg:
+        return True
+
+    info = await asyncio.create_subprocess_shell(f'jpeginfo -c {filename}',
+                                                 stdout=asyncio.subprocess.PIPE,
+                                                 stderr=asyncio.subprocess.PIPE)
+    out, _ = await info.communicate()
+    return "OK" in out.decode()
+
+
+async def delete_corrupt_images(image_folder: str, check_jpeg: bool = False) -> None:
+    logging.info('deleting corrupt images')
+    n_deleted = 0
+    for image in glob(f'{image_folder}/*.jpg'):
+        if not await is_valid_image(image, check_jpeg):
+            logging.debug(f'  deleting image {image}')
+            os.remove(image)
+            n_deleted += 1
+
+    logging.info(f'deleted {n_deleted} images')
+
+
 def create_resource_paths(organization_name: str, project_name: str, image_ids: List[str]) -> Tuple[List[str], List[str]]:
-    # TODO: experimental:
     return [f'/{organization_name}/projects/{project_name}/images/{id}/main' for id in image_ids], image_ids
-    # if not image_ids:
-    #     return [], []
-    # url_ids: List[Tuple(str, str)] = [(f'/{organization_name}/projects/{project_name}/images/{id}/main', id)
-    #                                   for id in image_ids]
-    # urls, ids = list(map(list, zip(*url_ids)))
-
-    # return urls, ids
 
 
 def create_image_folder(project_folder: str) -> str:
@@ -74,6 +98,24 @@ def create_image_folder(project_folder: str) -> str:
     return image_folder
 
 
+def read_or_create_uuid(identifier: str) -> str:
+    identifier = identifier.lower().replace(' ', '_')
+    uuids = {}
+    os.makedirs(GLOBALS.data_folder, exist_ok=True)
+    file_path = f'{GLOBALS.data_folder}/uuids.json'
+    if os.path.exists(file_path):
+        with open(file_path, 'r') as f:
+            uuids = json.load(f)
+
+    uuid = uuids.get(identifier, None)
+    if not uuid:
+        uuid = str(uuid4())
+        uuids[identifier] = uuid
+        with open(file_path, 'w') as f:
+            json.dump(uuids, f)
+    return uuid
+
+
 def ensure_socket_response(func):
     """Decorator to ensure that the return value of a socket.io event handler is a SocketResponse.
 
@@ -90,20 +132,85 @@ def ensure_socket_response(func):
 
             if isinstance(value, str):
                 return asdict(SocketResponse.for_success(value))
-            elif isinstance(value, bool):
+            if isinstance(value, bool):
                 return asdict(SocketResponse.from_bool(value))
-            elif isinstance(value, SocketResponse):
+            if isinstance(value, SocketResponse):
                 return value
-            elif (args[0] in ['connect', 'disconnect', 'connect_error']):
+            if (args[0] in ['connect', 'disconnect', 'connect_error']):
                 return value
-            elif value is None:
+            if value is None:
                 return None
-            else:
-                raise Exception(
-                    f"Return type for sio must be str, bool, SocketResponse or None', but was {type(value)}'")
+
+            raise Exception(
+                f"Return type for sio must be str, bool, SocketResponse or None', but was {type(value)}'")
         except Exception as e:
             logging.exception(f'An error occured for {args[0]}')
 
             return asdict(SocketResponse.for_failure(str(e)))
 
     return wrapper_ensure_socket_response
+
+
+def is_valid_uuid4(val):
+    if not val:
+        return False
+    try:
+        _ = UUID(str(val)).version
+        return True
+    except ValueError:
+        return False
+
+
+def create_project_folder(context: Context) -> str:
+    project_folder = f'{GLOBALS.data_folder}/{context.organization}/{context.project}'
+    os.makedirs(project_folder, exist_ok=True)
+    return project_folder
+
+
+def activate_asyncio_warnings() -> None:
+    '''Produce warnings for coroutines which take too long on the main loop and hence clog the event loop'''
+    try:
+        if sys.version_info.major >= 3 and sys.version_info.minor >= 7:  # most
+            loop = asyncio.get_running_loop()
+        else:
+            loop = asyncio.get_event_loop()
+
+        loop.set_debug(True)
+        loop.slow_callback_duration = 0.2
+        logging.info('activated asyncio warnings')
+    except Exception:
+        logging.exception('could not activate asyncio warnings. Exception:')
+
+
+def images_for_ids(image_ids, image_folder) -> List[str]:
+    logging.info(f'### Going to get images for {len(image_ids)} images ids')
+    start = perf_counter()
+    images = [img for img in glob(f'{image_folder}/**/*.*', recursive=True)
+              if os.path.splitext(os.path.basename(img))[0] in image_ids]
+    end = perf_counter()
+    logging.info(f'found {len(images)} images for {len(image_ids)} image ids, which took {end-start:0.2f} seconds')
+    return images
+
+
+def generate_training(project_folder: str, context: Context) -> Training:
+    training_uuid = str(uuid4())
+    return Training(
+        id=training_uuid,
+        context=context,
+        project_folder=project_folder,
+        images_folder=create_image_folder(project_folder),
+        training_folder=create_training_folder(project_folder, training_uuid)
+    )
+
+
+def delete_all_training_folders(project_folder: str):
+    if not os.path.exists(f'{project_folder}/trainings'):
+        return
+    for uuid in os.listdir(f'{project_folder}/trainings'):
+        shutil.rmtree(f'{project_folder}/trainings/{uuid}', ignore_errors=True)
+
+
+def create_training_folder(project_folder: str, trainings_id: str) -> str:
+    training_folder = f'{project_folder}/trainings/{trainings_id}'
+    os.makedirs(training_folder, exist_ok=True)
+    return training_folder

learning_loop_node/loop_communication.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import List, Optional
+from typing import Awaitable, Callable, List, Optional
 
 import httpx
 from httpx import Cookies, Timeout
@@ -24,21 +24,21 @@ class LoopCommunicator():
         self.project: str = environment_reader.project()  # used by mock_detector
         self.base_url: str = f'http{"s" if "learning-loop.ai" in host else ""}://' + host
         self.async_client: httpx.AsyncClient = httpx.AsyncClient(base_url=self.base_url, timeout=Timeout(60.0))
+        self.async_client.cookies.clear()
 
         logging.info(f'Loop interface initialized with base_url: {self.base_url} / user: {self.username}')
 
-    # @property
-    # def project_path(self):  # TODO: remove?
-    #     return f'/{self.organization}/projects/{self.project}'
+    def websocket_url(self) -> str:
+        return f'ws{"s" if "learning-loop.ai" in self.host else ""}://' + self.host
 
-    async def ensure_login(self) -> None:
+    async def ensure_login(self, relogin=False) -> None:
         """aiohttp client session needs to be created on the event loop"""
 
         assert not self.async_client.is_closed, 'async client must not be used after shutdown'
-        if not self.async_client.cookies.keys():
+        if not self.async_client.cookies.keys() or relogin:
+            self.async_client.cookies.clear()
             response = await self.async_client.post('/api/login', data={'username': self.username, 'password': self.password})
             if response.status_code != 200:
-                self.async_client.cookies.clear()
                 logging.info(f'Login failed with response: {response}')
                 raise LoopCommunicationException('Login failed with response: ' + str(response))
             self.async_client.cookies.update(response.cookies)
@@ -50,8 +50,9 @@ class LoopCommunicator():
         if response.status_code != 200:
             logging.info(f'Logout failed with response: {response}')
             raise LoopCommunicationException('Logout failed with response: ' + str(response))
+        self.async_client.cookies.clear()
 
-    async def get_cookies(self) -> Cookies:
+    def get_cookies(self) -> Cookies:
         return self.async_client.cookies
 
     async def shutdown(self):
@@ -70,37 +71,68 @@ class LoopCommunicator():
                 logging.info(f'backend not ready: {e}')
             await asyncio.sleep(10)
 
+    async def retry_on_401(self, func: Callable[..., Awaitable[httpx.Response]], *args, **kwargs) -> httpx.Response:
+        response = await func(*args, **kwargs)
+        if response.status_code == 401:
+            await self.ensure_login(relogin=True)
+            response = await func(*args, **kwargs)
+        return response
+
     async def get(self, path: str, requires_login: bool = True, api_prefix: str = '/api') -> httpx.Response:
         if requires_login:
             await self.ensure_login()
+            return await self.retry_on_401(self._get, path, api_prefix)
+        else:
+            return await self._get(path, api_prefix)
+
+    async def _get(self, path: str, api_prefix: str) -> httpx.Response:
         return await self.async_client.get(api_prefix+path)
 
-    async def put(self, path, files: Optional[List[str]]=None, requires_login=True, api_prefix='/api', **kwargs) -> httpx.Response:
+    async def put(self, path: str, files: Optional[List[str]] = None, requires_login: bool = True, api_prefix: str = '/api', **kwargs) -> httpx.Response:
         if requires_login:
             await self.ensure_login()
+            return await self.retry_on_401(self._put, path, files, api_prefix, **kwargs)
+        else:
+            return await self._put(path, files, api_prefix, **kwargs)
+
+    async def _put(self, path: str, files: Optional[List[str]], api_prefix: str, **kwargs) -> httpx.Response:
         if files is None:
             return await self.async_client.put(api_prefix+path, **kwargs)
-        
-        file_list = [('files', open(f, 'rb')) for f in files]  # TODO: does this properly close the files after upload?
-        return await self.async_client.put(api_prefix+path, files=file_list)
 
-    async def post(self, path, requires_login=True, api_prefix='/api', **kwargs) -> httpx.Response:
+        file_handles = []
+        for f in files:
+            try:
+                file_handles.append(open(f, 'rb'))
+            except FileNotFoundError:
+                for fh in file_handles:
+                    fh.close()  # Ensure all files are closed
+                return httpx.Response(404, content=b'File not found')
+
+        try:
+            file_list = [('files', fh) for fh in file_handles]  # Use file handles
+            response = await self.async_client.put(api_prefix+path, files=file_list)
+        finally:
+            for fh in file_handles:
+                fh.close()  # Ensure all files are closed
+
+        return response
+
+    async def post(self, path: str, requires_login: bool = True, api_prefix: str = '/api', **kwargs) -> httpx.Response:
         if requires_login:
             await self.ensure_login()
+            return await self.retry_on_401(self._post, path, api_prefix, **kwargs)
+        else:
+            return await self._post(path, api_prefix, **kwargs)
+
+    async def _post(self, path, api_prefix='/api', **kwargs) -> httpx.Response:
         return await self.async_client.post(api_prefix+path, **kwargs)
 
-    async def delete(self, path, requires_login=True, api_prefix='/api', **kwargs) -> httpx.Response:
+    async def delete(self, path: str, requires_login: bool = True, api_prefix: str = '/api', **kwargs) -> httpx.Response:
         if requires_login:
             await self.ensure_login()
-        return await self.async_client.delete(api_prefix+path, **kwargs)
-
-    # --------------------------------- unused?! --------------------------------- #TODO remove?
+            return await self.retry_on_401(self._delete, path, api_prefix, **kwargs)
+        else:
+            return await self._delete(path, api_prefix, **kwargs)
 
-    # def get_data(self, path):
-    #     return asyncio.get_event_loop().run_until_complete(self._get_data_async(path))
-
-    # async def _get_data_async(self, path) -> bytes:
-    #     response = await self.get(f'{self.project_path}{path}')
-    #     if response.status_code != 200:
-    #         raise LoopCommunicationException('bad response: ' + str(response))
-    #     return response.content
+    async def _delete(self, path, api_prefix, **kwargs) -> httpx.Response:
+        return await self.async_client.delete(api_prefix+path, **kwargs)

learning_loop_node/node.py

@@ -1,58 +1,58 @@
 import asyncio
-import json
 import logging
-import os
 import sys
 from abc import abstractmethod
+from contextlib import asynccontextmanager
 from datetime import datetime
-from typing import Optional
-from uuid import uuid4
+from typing import Any, Optional
 
 import aiohttp
 import socketio
 from fastapi import FastAPI
-from fastapi_utils.tasks import repeat_every
 from socketio import AsyncClient
 
-from .data_classes import Context, NodeState, NodeStatus
+from .data_classes import NodeStatus
 from .data_exchanger import DataExchanger
-from .globals import GLOBALS
-from .helpers import environment_reader, log_conf
-from .helpers.misc import ensure_socket_response
+from .helpers import log_conf
+from .helpers.misc import activate_asyncio_warnings, ensure_socket_response, read_or_create_uuid
 from .loop_communication import LoopCommunicator
 
 
 class Node(FastAPI):
 
-    def __init__(self, name: str, uuid: Optional[str] = None):
+    def __init__(self, name: str, uuid: Optional[str] = None, node_type: str = 'node', needs_login: bool = True):
         """Base class for all nodes. A node is a process that communicates with the zauberzeug learning loop.
+        This class provides the basic functionality to connect to the learning loop via socket.io and to exchange data.
 
         Args:
             name (str): The name of the node. This name is used to generate a uuid.
             uuid (Optional[str]): The uuid of the node. If None, a uuid is generated based on the name 
                 and stored in f'{GLOBALS.data_folder}/uuids.json'. 
-                From the second run, the uuid is recovered based on the name of the node. Defaults to None.
+                From the second run, the uuid is recovered based on the name of the node.
+            needs_login (bool): If True, the node will try to login to the learning loop.
         """
 
-        super().__init__()
+        super().__init__(lifespan=self.lifespan)
         log_conf.init()
 
+        self.name = name
+        self.uuid = uuid or read_or_create_uuid(self.name)
+        self.needs_login = needs_login
+
         self.log = logging.getLogger()
         self.loop_communicator = LoopCommunicator()
+        self.websocket_url = self.loop_communicator.websocket_url()
         self.data_exchanger = DataExchanger(None, self.loop_communicator)
 
-        host = environment_reader.host(default='learning-loop.ai')
-        self.ws_url = f'ws{"s" if "learning-loop.ai" in host else ""}://' + host
-
-        self.name = name
-        self.uuid = self.read_or_create_uuid(self.name) if uuid is None else uuid
-        self.startup_time = datetime.now()
+        self.startup_datetime = datetime.now()
         self._sio_client: Optional[AsyncClient] = None
         self.status = NodeStatus(id=self.uuid, name=self.name)
-        # NOTE this is can be set to False for Nodes which do not need to authenticate with the backend (like the DetectorNode)
-        self.needs_login = True
-        self._setup_sio_headers()
-        self._register_lifecycle_events()
+
+        self.sio_headers = {'organization': self.loop_communicator.organization,
+                            'project': self.loop_communicator.project,
+                            'nodeType': node_type}
+
+        self.repeat_task: Any = None
 
     @property
     def sio_client(self) -> AsyncClient:
@@ -60,52 +60,25 @@ class Node(FastAPI):
             raise Exception('sio_client not yet initialized')
         return self._sio_client
 
-    def sio_is_initialized(self) -> bool:
-        return self._sio_client is not None
-
-    # --------------------------------------------------- INIT ---------------------------------------------------
-
-    def read_or_create_uuid(self, identifier: str) -> str:
-        identifier = identifier.lower().replace(' ', '_')
-        uuids = {}
-        os.makedirs(GLOBALS.data_folder, exist_ok=True)
-        file_path = f'{GLOBALS.data_folder}/uuids.json'
-        if os.path.exists(file_path):
-            with open(file_path, 'r') as f:
-                uuids = json.load(f)
-
-        uuid = uuids.get(identifier, None)
-        if not uuid:
-            uuid = str(uuid4())
-            uuids[identifier] = uuid
-            with open(file_path, 'w') as f:
-                json.dump(uuids, f)
-        return uuid
-
-    def _setup_sio_headers(self) -> None:
-        self.sio_headers = {'organization': self.loop_communicator.organization,
-                            'project': self.loop_communicator.project,
-                            'nodeType': self.get_node_type()}
-
     # --------------------------------------------------- APPLICATION LIFECYCLE ---------------------------------------------------
-
-    def _register_lifecycle_events(self):
-        @self.on_event("startup")
-        async def startup():
+    @asynccontextmanager
+    async def lifespan(self, app: FastAPI):  # pylint: disable=unused-argument
+        try:
             await self._on_startup()
-
-        @self.on_event("shutdown")  # NOTE only used for developent ?!
-        async def shutdown():
+            self.repeat_task = asyncio.create_task(self.repeat_loop())
+            yield
+        finally:
             await self._on_shutdown()
-
-        @self.on_event("startup")
-        @repeat_every(seconds=5, raise_exceptions=False, wait_first=False)
-        async def ensure_connected() -> None:
-            await self._on_repeat()
+            if self.repeat_task is not None:
+                self.repeat_task.cancel()
+                try:
+                    await self.repeat_task
+                except asyncio.CancelledError:
+                    pass
 
     async def _on_startup(self):
         self.log.info('received "startup" lifecycle-event')
-        Node._activate_asyncio_warnings()
+        # activate_asyncio_warnings()
         if self.needs_login:
             await self.loop_communicator.backend_ready()
             self.log.info('ensuring login')
@@ -123,10 +96,18 @@ class Node(FastAPI):
         self.log.info('successfully disconnected from loop.')
         await self.on_shutdown()
 
+    async def repeat_loop(self) -> None:
+        """NOTE: with the lifespan approach, we cannot use @repeat_every anymore :("""
+        while True:
+            try:
+                await self._on_repeat()
+            except asyncio.CancelledError:
+                return
+            except Exception as e:
+                self.log.exception(f'error in repeat loop: {e}')
+            await asyncio.sleep(5)
+
     async def _on_repeat(self):
-        while not self.sio_is_initialized():
-            self.log.info('Waiting for sio client to be initialized')
-            await asyncio.sleep(1)
         if not self.sio_client.connected:
             self.log.info('Reconnecting to loop via sio')
             await self.connect_sio()
@@ -138,8 +119,11 @@ class Node(FastAPI):
     # --------------------------------------------------- SOCKET.IO ---------------------------------------------------
 
     async def create_sio_client(self):
-        cookies = await self.loop_communicator.get_cookies()
-        self._sio_client = AsyncClient(request_timeout=20, http_session=aiohttp.ClientSession(cookies=cookies))
+        """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."""
+
+        self._sio_client = AsyncClient(request_timeout=20,
+                                       http_session=aiohttp.ClientSession(cookies=self.loop_communicator.get_cookies()))
 
         # pylint: disable=protected-access
         self.sio_client._trigger_event = ensure_socket_response(self.sio_client._trigger_event)
@@ -147,72 +131,39 @@ class Node(FastAPI):
         @self._sio_client.event
         async def connect():
             self.log.info('received "connect" via sio from loop.')
-            self.status = NodeStatus(id=self.uuid, name=self.name)
-            state = await self.get_state()
-            try:
-                await self._update_send_state(state)
-            except:
-                self.log.exception('Error sending state. Exception:')
-                raise
 
         @self._sio_client.event
         async def disconnect():
             self.log.info('received "disconnect" via sio from loop.')
-            await self._update_send_state(NodeState.Offline)
 
         @self._sio_client.event
         async def restart():
-            self.log.info('received "restart" via sio from loop.')
-            self.restart()
+            self.log.info('received "restart" via sio from loop -> restarting node.')
+            sys.exit(0)
 
         self.register_sio_events(self._sio_client)
 
     async def connect_sio(self):
-        if not self.sio_is_initialized():
-            self.log.warning('sio client not yet initialized')
-            return
         try:
             await self.sio_client.disconnect()
         except Exception:
             pass
 
-        self.log.info(f'(re)connecting to Learning Loop at {self.ws_url}')
+        self.log.info(f'(re)connecting to Learning Loop at {self.websocket_url}')
         try:
-            await self.sio_client.connect(f"{self.ws_url}", headers=self.sio_headers, socketio_path="/ws/socket.io")
+            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.ws_url}". Exception:')
-
-    async def _update_send_state(self, state: NodeState):
-        self.status.state = state
-        if self.status.state != NodeState.Offline:
-            await self.send_status()
+            self.log.exception(f'error while connecting to "{self.websocket_url}". Exception:')
 
     # --------------------------------------------------- ABSTRACT METHODS ---------------------------------------------------
 
-    @abstractmethod
-    def register_sio_events(self, sio_client: AsyncClient):
-        """Register socket.io events for the communication with the learning loop.
-        The events: connect and disconnect are already registered and should not be overwritten."""
-
-    @abstractmethod
-    async def send_status(self):
-        """Send the current status to the learning loop.
-        Note that currently this method is also used to react to the response of the learning loop."""
-
-    @abstractmethod
-    async def get_state(self) -> NodeState:
-        """Return the current state of the node."""
-
-    @abstractmethod
-    def get_node_type(self):
-        pass
-
     @abstractmethod
     async def on_startup(self):
-        """This method is called when the node is started."""
+        """This method is called when the node is started.
+        Note: In this method the sio connection is not yet established!"""
 
     @abstractmethod
     async def on_shutdown(self):
@@ -221,32 +172,8 @@ class Node(FastAPI):
     @abstractmethod
     async def on_repeat(self):
         """This method is called every 10 seconds."""
-    # --------------------------------------------------- SHARED FUNCTIONS ---------------------------------------------------
-
-    def restart(self):
-        """Restart the node."""
-        self.log.info('restarting node')
-        sys.exit(0)
-
-    # --------------------------------------------------- HELPER ---------------------------------------------------
-
-    @staticmethod
-    def create_project_folder(context: Context) -> str:
-        project_folder = f'{GLOBALS.data_folder}/{context.organization}/{context.project}'
-        os.makedirs(project_folder, exist_ok=True)
-        return project_folder
 
-    @staticmethod
-    def _activate_asyncio_warnings() -> None:
-        '''Produce warnings for coroutines which take too long on the main loop and hence clog the event loop'''
-        try:
-            if sys.version_info.major >= 3 and sys.version_info.minor >= 7:  # most
-                loop = asyncio.get_running_loop()
-            else:
-                loop = asyncio.get_event_loop()
-
-            loop.set_debug(True)
-            loop.slow_callback_duration = 0.2
-            logging.info('activated asyncio warnings')
-        except Exception:
-            logging.exception('could not activate asyncio warnings. Exception:')
+    @abstractmethod
+    def register_sio_events(self, sio_client: AsyncClient):
+        """Register (additional) socket.io events for the communication with the learning loop.
+        The events: connect, disconnect and restart are already registered and should not be overwritten."""

learning_loop_node/tests/test_downloader.py

@@ -2,9 +2,10 @@ import os
 import shutil
 
 from learning_loop_node.data_classes import Context
-from learning_loop_node.data_exchanger import DataExchanger, check_jpeg
+from learning_loop_node.data_exchanger import DataExchanger
 from learning_loop_node.globals import GLOBALS
 
+from ..helpers.misc import delete_corrupt_images
 from . import test_helper
 
 
@@ -33,26 +34,26 @@ async def test_download_model(data_exchanger: DataExchanger):
 
 # pylint: disable=redefined-outer-name
 async def test_fetching_image_ids(data_exchanger: DataExchanger):
-    ids = await data_exchanger.fetch_image_ids()
+    ids = await data_exchanger.fetch_image_uuids()
     assert len(ids) == 3
 
 
 async def test_download_images(data_exchanger: DataExchanger):
     _, image_folder, _ = test_helper.create_needed_folders()
-    image_ids = await data_exchanger.fetch_image_ids()
+    image_ids = await data_exchanger.fetch_image_uuids()
     await data_exchanger.download_images(image_ids, image_folder)
     files = test_helper.get_files_in_folder(GLOBALS.data_folder)
     assert len(files) == 3
 
 
 async def test_download_training_data(data_exchanger: DataExchanger):
-    image_ids = await data_exchanger.fetch_image_ids()
+    image_ids = await data_exchanger.fetch_image_uuids()
     image_data = await data_exchanger.download_images_data(image_ids)
     assert len(image_data) == 3
 
 
 async def test_removal_of_corrupted_images(data_exchanger: DataExchanger):
-    image_ids = await data_exchanger.fetch_image_ids()
+    image_ids = await data_exchanger.fetch_image_uuids()
 
     shutil.rmtree('/tmp/img_folder', ignore_errors=True)
     os.makedirs('/tmp/img_folder', exist_ok=True)
@@ -65,7 +66,7 @@ async def test_removal_of_corrupted_images(data_exchanger: DataExchanger):
     with open('/tmp/img_folder/c1.jpg', 'w') as f:
         f.write('I am no image')
 
-    await data_exchanger.delete_corrupt_images('/tmp/img_folder')
+    await delete_corrupt_images('/tmp/img_folder', True)
 
-    assert len(os.listdir('/tmp/img_folder')) == num_images if check_jpeg else num_images - 1
+    assert len(os.listdir('/tmp/img_folder')) == num_images if data_exchanger.check_jpeg else num_images - 1
     shutil.rmtree('/tmp/img_folder', ignore_errors=True)