mercuto-client 0.3.0__py3-none-any.whl → 0.3.4a3__py3-none-any.whl

mercuto_client/_tests/test_ingester/test_file_processor.py

@@ -2,7 +2,7 @@ import os
 import sqlite3
 import tempfile
 import time
-from typing import Generator, Tuple
+from typing import Generator, Iterator, Tuple
 
 import pytest
 
@@ -14,30 +14,40 @@ def mock_process_callback(filepath: str) -> bool:
 
 
 @pytest.fixture
-def temp_env() -> Generator[Tuple[FileProcessor, str, str], None, None]:
+def buffer_directory() -> Iterator[str]:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        yield temp_dir
+
+
+@pytest.fixture
+def work_directory() -> Iterator[str]:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        yield temp_dir
+
+
+@pytest.fixture
+def database_path(work_directory: str) -> Iterator[str]:
+    yield os.path.join(work_directory, "test_buffer.db")
+
+
+@pytest.fixture
+def temp_env(buffer_directory: str, database_path: str) -> Generator[Tuple[FileProcessor, str, str], None, None]:
     """Setup temporary directory and database"""
-    buffer_dir: str = tempfile.mkdtemp()
-    workdir = tempfile.mkdtemp()
-    db_path: str = os.path.join(workdir, "test_buffer.db")
 
     processor: FileProcessor = FileProcessor(
-        buffer_dir=buffer_dir,
-        db_path=db_path,
+        buffer_dir=buffer_directory,
+        db_path=database_path,
         max_files=3,
         max_attempts=2,
         process_callback=mock_process_callback
     )
 
-    yield processor, buffer_dir, db_path
-
-    # Cleanup
-    if os.path.exists(db_path):
-        os.remove(db_path)
+    yield processor, buffer_directory, database_path
 
 
 def test_init_db(temp_env: Tuple[FileProcessor, str, str]) -> None:
     """Verify database initialization"""
-    processor, _, db_path = temp_env
+    _, _, db_path = temp_env
     conn: sqlite3.Connection = sqlite3.connect(db_path)
     cursor: sqlite3.Cursor = conn.cursor()
     cursor.execute(
@@ -208,3 +218,151 @@ def test_scan_existing_files_that_havnt_been_processed(temp_env: Tuple[FileProce
     files: list[tuple[str]] = cursor.fetchall()
     conn.close()
     assert files == [('file0.txt',), ('file1.txt',), ('file2.txt',), ('file3.txt',), ('file4.txt',)]
+
+
+def test_ensure_free_space(database_path: str, buffer_directory: str) -> None:
+    """Ensure that keeping X MB free space works correctly"""
+    remaining_free_space_mb = 501
+
+    def mock_free_space_checker(dir: str) -> float:
+        """
+        Mock free space checker that returns a controlled value.
+        Returns remaining_free_space_mb - number of files in the buffer directory.
+        """
+        n_files = len(os.listdir(buffer_directory))
+        return remaining_free_space_mb - n_files
+
+    processor = FileProcessor(
+        buffer_dir=buffer_directory,
+        db_path=database_path,
+        max_files=10,
+        max_attempts=2,
+        process_callback=mock_process_callback,
+        target_free_space_mb=500,
+        free_space_checker=mock_free_space_checker
+    )
+
+    def create_and_add_file(name: str) -> str:
+        file_path = os.path.join(buffer_directory, name)
+        with open(file_path, 'w') as f:
+            f.write("Test content")
+        processor.add_file_to_db(file_path)
+        return file_path
+
+    create_and_add_file('success_file_1.txt')
+    create_and_add_file('success_file_2.txt')
+    create_and_add_file('success_file_3.txt')
+    create_and_add_file('success_file_4.txt')
+
+    # There should be 4 files on disk
+    assert len(os.listdir(buffer_directory)) == 4
+
+    # Trigger cleanup, should not delete anything as we have enough free space
+    remaining_free_space_mb = 504
+    processor.cleanup_old_files()
+    assert len(os.listdir(buffer_directory)) == 4
+
+    # Reduce free space to trigger cleanup
+    remaining_free_space_mb = 502
+    processor.cleanup_old_files()
+    # This should have deleted the oldest two files to maintain 500 MB free space
+    # E.g. no deletes gives 498 MB free space, deleting one file gives 499 MB free space,
+    # deleting two files gives 500 MB free space
+    assert len(os.listdir(buffer_directory)) == 2
+    remaining_files = os.listdir(buffer_directory)
+    assert set(remaining_files) == {'success_file_3.txt', 'success_file_4.txt'}
+
+
+def test_ensure_free_space_with_externally_modified_files(database_path: str, buffer_directory: str) -> None:
+    """
+    Ensure that keeping X MB free space works correctly under the following:
+    - Files may be deleted from disk externally (not via the processor)
+    """
+    remaining_free_space_mb = 501
+
+    def mock_free_space_checker(dir: str) -> float:
+        """
+        Mock free space checker that returns a controlled value.
+        Returns remaining_free_space_mb - number of files in the buffer directory.
+        """
+        n_files = len(os.listdir(buffer_directory))
+        return remaining_free_space_mb - n_files
+
+    processor = FileProcessor(
+        buffer_dir=buffer_directory,
+        db_path=database_path,
+        max_files=10,
+        max_attempts=2,
+        process_callback=mock_process_callback,
+        target_free_space_mb=500,
+        free_space_checker=mock_free_space_checker
+    )
+
+    def create_and_add_file(name: str) -> str:
+        file_path = os.path.join(buffer_directory, name)
+        with open(file_path, 'w') as f:
+            f.write("Test content")
+        processor.add_file_to_db(file_path)
+        return file_path
+
+    file1 = create_and_add_file('success_file_1.txt')
+    file2 = create_and_add_file('success_file_2.txt')
+    create_and_add_file('success_file_3.txt')
+    create_and_add_file('success_file_4.txt')
+
+    # There should be 4 files on disk
+    assert len(os.listdir(buffer_directory)) == 4
+
+    # Manually delete files 1 and 2 to simulate external deletion
+    os.remove(file1)
+    os.remove(file2)
+
+    # Reduce free space to trigger cleanup
+    remaining_free_space_mb = 100
+    processor.cleanup_old_files()
+    # This should still delete all files (even though 1 and 2 are already gone) to maintain 500 MB free space
+    # The current free space of 100MB will never be reached, so all files should be deleted.
+    assert len(os.listdir(buffer_directory)) == 0
+
+
+def test_cleanup_max_files_with_externally_modified_files(database_path: str, buffer_directory: str) -> None:
+    """
+    Ensure that keeping X max files works correctly under the following:
+    - Files may be deleted from disk externally (not via the processor)
+    """
+
+    processor = FileProcessor(
+        buffer_dir=buffer_directory,
+        db_path=database_path,
+        max_files=1,
+        max_attempts=2,
+        process_callback=mock_process_callback,
+        target_free_space_mb=None,
+    )
+
+    def create_and_add_file(name: str) -> str:
+        file_path = os.path.join(buffer_directory, name)
+        with open(file_path, 'w') as f:
+            f.write("Test content")
+        processor.add_file_to_db(file_path)
+        return file_path
+
+    file1 = create_and_add_file('success_file_1.txt')
+    file2 = create_and_add_file('success_file_2.txt')
+    create_and_add_file('success_file_3.txt')
+    create_and_add_file('success_file_4.txt')
+
+    # There should be 4 files on disk
+    assert len(os.listdir(buffer_directory)) == 4
+
+    # Manually delete files 1 and 2 to simulate external deletion
+    os.remove(file1)
+    os.remove(file2)
+
+    # Trigger cleanup
+    processor.cleanup_old_files()
+    # THis should have deleted all but one file to maintain max_files=1
+    assert len(os.listdir(buffer_directory)) == 1
+    # It should be the last file added
+    remaining_files = os.listdir(buffer_directory)
+    assert set(remaining_files) == {'success_file_4.txt'}

mercuto_client/acl.py

@@ -20,6 +20,13 @@ class ResourceTypes:
         TENANT = 'tenant'
         WILDCARD = '*'
 
+    class Notifications:
+        """
+        Resource types available for ServiceTypes.NOTIFICATIONS
+        """
+        CONTACT_GROUP = 'contact-group'
+        WILDCARD = '*'
+
 
 class AllowedActions:
     WILDCARD = '*'
@@ -46,6 +53,13 @@ class AllowedActions:
         VIEW_USER_DETAILED_INFO = 'IDENTITY:VIEW_USER_DETAILED_INFO'
         CREATE_NEW_TENANTS = 'IDENTITY:CREATE_NEW_TENANTS'
 
+    class Notifications:
+        """
+        Actions available for ServiceTypes.NOTIFICATIONS
+        """
+        WILDCARD = '*'
+        ISSUE_NOTIFICATION = 'NOTIFICATIONS:ISSUE_NOTIFICATION'
+
 
 class ServiceTypes:
     class Mercuto:
@@ -58,8 +72,14 @@ class ServiceTypes:
         ResourceTypes = ResourceTypes.Identity
         AllowedActions = AllowedActions.Identity
 
+    class Notifications:
+        Name = 'notifications'
+        ResourceTypes = ResourceTypes.Notifications
+        AllowedActions = AllowedActions.Notifications
+
     IDENTITY = Identity.Name
     MERCUTO = Mercuto.Name
+    NOTIFICATIONS = Notifications.Name
     WILDCARD = '*'
 
 

mercuto_client/client.py

@@ -3,7 +3,8 @@ import json as json_stdlib
 import logging
 import os
 import time
-from typing import Any, Iterator, Mapping, Optional, Protocol, Type, TypeVar
+from typing import (Any, Iterator, Literal, Mapping, Optional, Protocol, Type,
+                    TypeVar)
 
 import requests
 import requests.cookies
@@ -15,6 +16,9 @@ from .modules.core import MercutoCoreService
 from .modules.data import MercutoDataService
 from .modules.fatigue import MercutoFatigueService
 from .modules.identity import MercutoIdentityService
+from .modules.media import MercutoMediaService
+from .modules.notifications import MercutoNotificationService
+from .modules.reports import MercutoReportService
 
 logger = logging.getLogger(__name__)
 
@@ -36,11 +40,11 @@ class MercutoClient:
         if url.endswith('/'):
             url = url[:-1]
 
-        if not url.startswith('https://'):
+        if verify_ssl and not url.startswith('https://'):
             raise ValueError(f'Url must be https, is {url}')
 
         self._url = url
-        self._verify_ssl = verify_ssl
+        self.verify_ssl = verify_ssl
 
         if active_session is None:
             self._current_session = requests.Session()
@@ -63,11 +67,11 @@ class MercutoClient:
             raise MercutoClientException("No credentials set")
         return self._auth_method.unique_key()
 
-    def set_verify_ssl(self, verify_ssl: bool) -> None:
-        self._verify_ssl = verify_ssl
+    def setverify_ssl(self, verify_ssl: bool) -> None:
+        self.verify_ssl = verify_ssl
 
     def copy(self) -> 'MercutoClient':
-        return MercutoClient(self._url, self._verify_ssl, self._current_session)
+        return MercutoClient(self._url, self.verify_ssl, self._current_session)
 
     @contextlib.contextmanager
     def as_credentials(self, api_key: Optional[str] = None,
@@ -78,7 +82,7 @@ class MercutoClient:
         Same as .connect(), but as a context manager. Will automatically logout when exiting the context.
         """
         # TODO: We are passing the current session along to re-use connections for speed. Will this cause security issues?
-        other = MercutoClient(self._url, self._verify_ssl, self._current_session)
+        other = MercutoClient(self._url, self.verify_ssl, self._current_session)
         try:
             yield other.connect(api_key=api_key, service_token=service_token, bearer_token=bearer_token, headers=headers)
         finally:
@@ -108,6 +112,26 @@ class MercutoClient:
         base.update(headers)
         return base
 
+    def session(self) -> requests.Session:
+        return self._current_session
+
+    def request(self, url: str, method: Literal['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'HEAD', 'OPTIONS'],
+                params: Optional[dict[str, Any]] = None,
+                json: Optional[dict[str, Any]] = None,
+                raise_for_status: bool = True,
+                **kwargs: Any) -> requests.Response:
+        """
+        Make an HTTP request to the Mercuto API.
+        :param url: The URL path (relative to the base API URL) to make the request to.
+        :param method: The HTTP method to use (e.g., 'GET', 'POST', etc.).
+        :param params: Optional dictionary of query parameters to include in the request.
+        :param json: Optional dictionary to send as a JSON payload in the request body.
+        :param raise_for_status: Whether to raise an exception for HTTP error responses.
+        :param kwargs: Additional keyword arguments to pass to the requests method.
+        :return: The HTTP response object.
+        """
+        return self._http_request(url, method, params=params, json=json, raise_for_status=raise_for_status, **kwargs)
+
     def _http_request(self, url: str, method: str,
                       params: Optional[dict[str, Any]] = None,
                       json: Optional[dict[str, Any]] = None,
@@ -122,7 +146,7 @@ class MercutoClient:
         kwargs['headers'] = self._update_headers(kwargs.get('headers', {}))
 
         if 'verify' not in kwargs:
-            kwargs['verify'] = self._verify_ssl
+            kwargs['verify'] = self.verify_ssl
 
         if 'cookies' not in kwargs:
             kwargs['cookies'] = self._cookies
@@ -167,6 +191,15 @@ class MercutoClient:
     def core(self) -> 'MercutoCoreService':
         return self._add_and_fetch_module('core', MercutoCoreService)
 
+    def media(self) -> 'MercutoMediaService':
+        return self._add_and_fetch_module('media', MercutoMediaService)
+
+    def reports(self) -> 'MercutoReportService':
+        return self._add_and_fetch_module('reports', MercutoReportService)
+
+    def notifications(self) -> 'MercutoNotificationService':
+        return self._add_and_fetch_module('notifications', MercutoNotificationService)
+
     def login(self, authentication: IAuthenticationMethod) -> None:
         self._auth_method = authentication
 

mercuto_client/ingester/__main__.py

@@ -40,9 +40,11 @@ if __name__ == '__main__':
                         help='Enable verbose output')
     parser.add_argument('-d', '--directory', type=str,
                         help='Directory to store ingested files. Default is a directory called `buffered-files` in the workdir.')
-    parser.add_argument('-s', '--size', type=int,
-                        help='Size in MB for total amount of files to store in the buffer. \
-                            Default is 75% of the available disk space on the buffer partition excluding the directory itself', default=None)
+    parser.add_argument('-s', '--target-free-space-mb', type=int,
+                        help='Size in MB for total amount of remaining free space to keep available. \
+                            Default is 25% of the available disk space on the buffer partition excluding the directory itself', default=None)
+    parser.add_argument('--max-files', type=int,
+                        help='Maximum number of files to keep in the buffer. Default is to use the size param.', default=None)
     parser.add_argument('--max-attempts', type=int,
                         help='Maximum number of attempts to process a file before giving up. Default is 1000.',
                         default=1000)
@@ -113,10 +115,10 @@ if __name__ == '__main__':
     ftp_dir = os.path.join(workdir, 'temp-ftp-data')
     os.makedirs(ftp_dir, exist_ok=True)
 
-    size = args.size
-    if size is None:
-        size = get_free_space_excluding_files(buffer_directory) * 0.75 // (1024 * 1024)  # Convert to MB
-        logging.info(f"Buffer size set to {size} MB based on available disk space.")
+    target_free_space_mb = args.target_free_space_mb
+    if target_free_space_mb is None and args.max_files is None:
+        target_free_space_mb = get_free_space_excluding_files(buffer_directory) * 0.25 // (1024 * 1024)  # Convert to MB
+        logging.info(f"Target remaining free space set to {target_free_space_mb} MB based on available disk space.")
 
     if args.mapping is not None:
         import json
@@ -146,7 +148,8 @@ if __name__ == '__main__':
         db_path=database_path,
         process_callback=ingester.process_file,
         max_attempts=args.max_attempts,
-        free_space_mb=size)
+        target_free_space_mb=target_free_space_mb,
+        max_files=args.max_files)
 
     processor.scan_existing_files()
 

mercuto_client/ingester/processor.py

@@ -8,6 +8,22 @@ from typing import Callable, Optional
 logger = logging.getLogger(__name__)
 
 
+def _default_standard_clock(_: str):
+    # Default clock function to get current time in seconds since epoch
+    return datetime.now(timezone.utc).timestamp()
+
+
+def _default_file_clock(filepath: str) -> float:
+    # Default clock function to get file creation time in seconds since epoch
+    return os.path.getctime(filepath)
+
+
+def _default_free_space_checker(buffer_dir: str) -> float:
+    """Returns the free space in MB on the partition where the buffer directory is located."""
+    _, _, free = shutil.disk_usage(buffer_dir)
+    return free / (1024 * 1024)
+
+
 class FileProcessor:
     """
     System for processing files in a strict order with retry logic.
@@ -20,11 +36,14 @@ class FileProcessor:
     :param process_callback: Callable that processes a file. Should return True if processing is successful
     :param max_attempts: Maximum number of attempts to process a file before marking it as failed.
     :param max_files: Maximum number of files to keep in the buffer directory. If None, no limit is enforced.
-    :param free_space_mb: Optional minimum free space in MB to keep on the partition where the buffer directory is located.
+    :param target_free_space_mb: Optional minimum free space in MB to keep on the partition where the buffer directory is located.
         This is combined with max_files to determine when to delete old files and takes precedence over max_files.
     :param clock: Optional callable that returns the timestamp for the file based on the filename. Takes in the file name as an argument
         and should return a float representing the timestamp in seconds since the epoch. Defaults to datetime.now().timestamp().
         This clock is NOT used when scanning existing files, it takes its own clock function.
+    :param free_space_checker: Optional callable that returns the free space in MB on the partition where the buffer directory is located.
+        Takes in the buffer directory path as an argument and should return a float representing the free space in MB.
+        Defaults to checking the actual free space on the disk.
 
 
     Provides a callback for processing files, which should return True if successful.
@@ -50,20 +69,18 @@ class FileProcessor:
                  process_callback: Callable[[str], bool],
                  max_attempts: int,
                  max_files: Optional[int] = None,
-                 free_space_mb: Optional[float] = None,
-                 clock: Optional[Callable[[str], float]] = None
+                 target_free_space_mb: Optional[float] = None,
+                 clock: Optional[Callable[[str], float]] = None,
+                 free_space_checker: Optional[Callable[[str], float]] = None
                  ) -> None:
         self._buffer_dir = buffer_dir
         self._db_path = db_path
         self._max_files = max_files
         self._max_attempts = max_attempts
         self._process_callback = process_callback
-        self._free_space_mb = free_space_mb
-        if clock is None:
-            # Default clock function to get current time in seconds since epoch
-            def clock(_): return datetime.now(timezone.utc).timestamp()
-        self._clock = clock
-
+        self._target_free_space_mb = target_free_space_mb
+        self._clock = clock if clock is not None else _default_standard_clock
+        self._free_space_checker = free_space_checker if free_space_checker is not None else _default_free_space_checker
         os.makedirs(self._buffer_dir, exist_ok=True)
         self._init_db()
 
@@ -102,9 +119,7 @@ class FileProcessor:
 
         """
         if clock is None:
-            # Default clock function to get current time in seconds since epoch
-            def clock(filepath: str) -> float:
-                return os.path.getctime(filepath)
+            clock = _default_file_clock
 
         conn = sqlite3.connect(self._db_path)
         cursor = conn.cursor()
@@ -216,38 +231,42 @@ class FileProcessor:
         conn.close()
 
         for (filepath,) in reversed(files_to_delete):
-            if os.path.exists(filepath):
+            conn = sqlite3.connect(self._db_path)
+            cursor = conn.cursor()
+            cursor.execute(
+                "DELETE FROM file_buffer WHERE filepath = ?", (filepath,))
+            conn.commit()
+            conn.close()
+            try:
                 os.remove(filepath)
-                conn = sqlite3.connect(self._db_path)
-                cursor = conn.cursor()
-                cursor.execute(
-                    "DELETE FROM file_buffer WHERE filepath = ?", (filepath,))
-                conn.commit()
-                conn.close()
                 logger.info(f"Deleted old file {filepath}")
+            except FileNotFoundError:
+                logger.warning(f"File {filepath} does not exist for deletion.")
+            except Exception as e:
+                logger.error(f"Error deleting file {filepath}: {e}")
 
     def cleanup_old_files(self) -> None:
         """Remove old files based on max_files and free space."""
         if self._max_files is not None:
             self.cleanup_old_files_with_max_files()
 
-        if self._free_space_mb is not None:
+        if self._target_free_space_mb is not None:
             self.cleanup_old_files_with_free_space()
 
     def cleanup_old_files_with_free_space(self) -> None:
         """Remove old files ensuring free space is maintained."""
-        if self._free_space_mb is None:
+        if self._target_free_space_mb is None:
             return
 
-        def free_space_mb() -> float:
-            """Returns the free space in MB on the partition where the buffer directory is located."""
-            total, used, free = shutil.disk_usage(self._buffer_dir)
-            return free / (1024 * 1024)
-
-        while free_space_mb() < self._free_space_mb:
+        count = 0
+        while self._free_space_checker(self._buffer_dir) < self._target_free_space_mb:
             if not self._delete_oldest_file():
                 logger.warning("No more files to delete to free up space.")
                 break
+            count += 1
+            if count > 1000:
+                logger.error("Too many iterations trying to free space, aborting.")
+                break
 
     def _delete_oldest_file(self) -> bool:
         """Deletes the oldest file in the buffer directory."""
@@ -258,25 +277,26 @@ class FileProcessor:
         oldest_file: Optional[tuple[str]] = cursor.fetchone()
         if not oldest_file:
             logger.info("No files to delete.")
+            conn.close()
             return False
 
         filepath = oldest_file[0]
-        removed = False
-        if os.path.exists(filepath):
+        cursor.execute(
+            "DELETE FROM file_buffer WHERE filepath = ?", (filepath,))
+        conn.commit()
+        conn.close()
+        try:
             os.remove(filepath)
-            cursor.execute(
-                "DELETE FROM file_buffer WHERE filepath = ?", (filepath,))
-            conn.commit()
-            removed = True
             logger.info(f"Deleted oldest file {filepath}")
-        else:
+        except FileNotFoundError:
             logger.warning(f"Oldest file {filepath} does not exist.")
+        except Exception as e:
+            logger.error(f"Error deleting oldest file {filepath}: {e}")
 
-        conn.close()
-        return removed
+        return True
 
     def add_file_to_db(self, filepath: str) -> None:
-        """Adds a new file to database and triggers processing."""
+        """Adds a new file to database to be processed on the next call to process_next_file."""
         timestamp = self._clock(filepath)
         filename: str = os.path.basename(filepath)
         conn = sqlite3.connect(self._db_path)

mercuto_client/mocks/__init__.py

@@ -10,6 +10,9 @@ from ..client import MercutoClient
 def mock_mercuto(data: bool = True,
                  identity: bool = True,
                  fatigue: bool = True,
+                 core: bool = True,
+                 media: bool = True,
+                 notifications: bool = True,
                  verify_service_token: Optional[Callable[[str], VerifyMyPermissions]] = None) -> Iterator[None]:
     """
     While this context is active, all calls to MercutoClient will use mocked services.
@@ -23,12 +26,40 @@ def mock_mercuto(data: bool = True,
         if data:
             stack.enter_context(mock_data_module())
         if identity:
-            stack.enter_context(mock_identity_module(verify_service_token=verify_service_token))
+            stack.enter_context(mock_identity_module(
+                verify_service_token=verify_service_token))
         if fatigue:
             stack.enter_context(mock_fatigue_module())
+        if core:
+            stack.enter_context(mock_core_module())
+        if media:
+            stack.enter_context(mock_media_module())
+        if notifications:
+            stack.enter_context(mock_notifications_module())
         yield
 
 
+@contextlib.contextmanager
+def mock_core_module() -> Iterator[None]:
+    from .mock_core import MockMercutoCoreService
+    original = MercutoClient.core
+
+    _cache: Optional[MockMercutoCoreService] = None
+
+    def stub(self: MercutoClient) -> MockMercutoCoreService:
+        nonlocal _cache
+        if _cache is None:
+            _cache = MockMercutoCoreService(self)
+        _cache._client = self
+        return _cache
+
+    try:
+        setattr(MercutoClient, 'core', stub)
+        yield
+    finally:
+        setattr(MercutoClient, 'core', original)
+
+
 @contextlib.contextmanager
 def mock_data_module() -> Iterator[None]:
     from .mock_data import MockMercutoDataService
@@ -60,7 +91,8 @@ def mock_identity_module(verify_service_token: Optional[Callable[[str], VerifyMy
     def stub(self: MercutoClient) -> MockMercutoIdentityService:
         nonlocal _cache
         if _cache is None:
-            _cache = MockMercutoIdentityService(self, verify_service_token=verify_service_token)
+            _cache = MockMercutoIdentityService(
+                self, verify_service_token=verify_service_token)
         _cache._client = self
         return _cache
 
@@ -90,3 +122,45 @@ def mock_fatigue_module() -> Iterator[None]:
         yield
     finally:
         setattr(MercutoClient, 'fatigue', original)
+
+
+@contextlib.contextmanager
+def mock_media_module() -> Iterator[None]:
+    from .mock_media import MockMercutoMediaService
+    original = MercutoClient.media
+
+    _cache: Optional[MockMercutoMediaService] = None
+
+    def stub(self: MercutoClient) -> MockMercutoMediaService:
+        nonlocal _cache
+        if _cache is None:
+            _cache = MockMercutoMediaService(self)
+        _cache._client = self
+        return _cache
+
+    try:
+        setattr(MercutoClient, 'media', stub)
+        yield
+    finally:
+        setattr(MercutoClient, 'media', original)
+
+
+@contextlib.contextmanager
+def mock_notifications_module() -> Iterator[None]:
+    from .mock_notifications import MockMercutoNotificationService
+    original = MercutoClient.notifications
+
+    _cache: Optional[MockMercutoNotificationService] = None
+
+    def stub(self: MercutoClient) -> MockMercutoNotificationService:
+        nonlocal _cache
+        if _cache is None:
+            _cache = MockMercutoNotificationService(self)
+        _cache._client = self
+        return _cache
+
+    try:
+        setattr(MercutoClient, 'notifications', stub)
+        yield
+    finally:
+        setattr(MercutoClient, 'notifications', original)