fix: Relax ServiceLocator restrictions (#837)

- this makes it possible to overwrite services in the service locator
before they are retrieved for the first time
- closes #806

Author: janbuchar

src/crawlee/_service_locator.py

@@ -20,15 +20,16 @@ def __init__(self) -> None:
         self._storage_client: BaseStorageClient | None = None
 
         # Flags to check if the services were already set.
-        self._configuration_was_set = False
-        self._event_manager_was_set = False
-        self._storage_client_was_set = False
+        self._configuration_was_retrieved = False
+        self._event_manager_was_retrieved = False
+        self._storage_client_was_retrieved = False
 
     def get_configuration(self) -> Configuration:
         """Get the configuration."""
         if self._configuration is None:
             self._configuration = Configuration()
 
+        self._configuration_was_retrieved = True
         return self._configuration
 
     def set_configuration(self, configuration: Configuration) -> None:
@@ -38,13 +39,12 @@ def set_configuration(self, configuration: Configuration) -> None:
             configuration: The configuration to set.
 
         Raises:
-            ServiceConflictError: If the configuration was already set.
+            ServiceConflictError: If the configuration has already been retrieved before.
         """
-        if self._configuration_was_set:
+        if self._configuration_was_retrieved:
             raise ServiceConflictError(Configuration, configuration, self._configuration)
 
         self._configuration = configuration
-        self._configuration_was_set = True
 
     def get_event_manager(self) -> EventManager:
         """Get the event manager."""
@@ -53,6 +53,7 @@ def get_event_manager(self) -> EventManager:
 
             self._event_manager = LocalEventManager()
 
+        self._event_manager_was_retrieved = True
         return self._event_manager
 
     def set_event_manager(self, event_manager: EventManager) -> None:
@@ -62,13 +63,12 @@ def set_event_manager(self, event_manager: EventManager) -> None:
             event_manager: The event manager to set.
 
         Raises:
-            ServiceConflictError: If the event manager was already set.
+            ServiceConflictError: If the event manager has already been retrieved before.
         """
-        if self._event_manager_was_set:
+        if self._event_manager_was_retrieved:
             raise ServiceConflictError(EventManager, event_manager, self._event_manager)
 
         self._event_manager = event_manager
-        self._event_manager_was_set = True
 
     def get_storage_client(self) -> BaseStorageClient:
         """Get the storage client."""
@@ -77,6 +77,7 @@ def get_storage_client(self) -> BaseStorageClient:
 
             self._storage_client = MemoryStorageClient.from_config()
 
+        self._storage_client_was_retrieved = True
         return self._storage_client
 
     def set_storage_client(self, storage_client: BaseStorageClient) -> None:
@@ -86,13 +87,12 @@ def set_storage_client(self, storage_client: BaseStorageClient) -> None:
             storage_client: The storage client to set.
 
         Raises:
-            ServiceConflictError: If the storage client was already set.
+            ServiceConflictError: If the storage client has already been retrieved before.
         """
-        if self._storage_client_was_set:
+        if self._storage_client_was_retrieved:
             raise ServiceConflictError(BaseStorageClient, storage_client, self._storage_client)
 
         self._storage_client = storage_client
-        self._storage_client_was_set = True
 
 
 service_locator = ServiceLocator()

src/crawlee/errors.py

@@ -38,11 +38,11 @@ class SessionError(Exception):
 
 @docs_group('Errors')
 class ServiceConflictError(Exception):
-    """Raised when attempting to reassign a service in service container that was already configured."""
+    """Raised when attempting to reassign a service in service container that is already in use."""
 
     def __init__(self, service: type, new_value: object, existing_value: object) -> None:
         super().__init__(
-            f'Service {service.__name__} has already been set. Existing value: {existing_value}, '
+            f'Service {service.__name__} is already in use. Existing value: {existing_value}, '
             f'attempted new value: {new_value}.'
         )
 

tests/unit/_memory_storage_client/test_memory_storage_client.py

@@ -230,7 +230,7 @@ async def test_not_implemented_method(tmp_path: Path) -> None:
 async def test_default_storage_path_used(monkeypatch: pytest.MonkeyPatch) -> None:
     # Reset the configuration in service locator
     service_locator._configuration = None
-    service_locator._configuration_was_set = False
+    service_locator._configuration_was_retrieved = False
 
     # Remove the env var for setting the storage directory
     monkeypatch.delenv('CRAWLEE_STORAGE_DIR', raising=False)

tests/unit/basic_crawler/test_basic_crawler.py

@@ -9,18 +9,21 @@
 from dataclasses import dataclass
 from datetime import timedelta
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 from unittest.mock import AsyncMock, Mock
 
 import httpx
 import pytest
 
-from crawlee import ConcurrencySettings, EnqueueStrategy, Glob
+from crawlee import ConcurrencySettings, EnqueueStrategy, Glob, service_locator
 from crawlee._request import BaseRequestData, Request
 from crawlee._types import BasicCrawlingContext, EnqueueLinksKwargs, HttpHeaders
 from crawlee.basic_crawler import BasicCrawler
 from crawlee.configuration import Configuration
 from crawlee.errors import SessionError, UserDefinedErrorHandlerError
+from crawlee.events._local_event_manager import LocalEventManager
+from crawlee.memory_storage_client import MemoryStorageClient
+from crawlee.memory_storage_client._dataset_client import DatasetClient
 from crawlee.request_loaders import RequestList, RequestManagerTandem
 from crawlee.statistics import FinalStatistics
 from crawlee.storages import Dataset, KeyValueStore, RequestQueue
@@ -997,3 +1000,49 @@ async def handler(context: BasicCrawlingContext) -> None:
     assert len(processed_urls) == 2
     assert stats.requests_total == 2
     assert stats.requests_finished == 2
+
+
+async def test_sets_services() -> None:
+    custom_configuration = Configuration()
+    custom_event_manager = LocalEventManager()
+    custom_storage_client = MemoryStorageClient.from_config(custom_configuration)
+
+    crawler = BasicCrawler(
+        configuration=custom_configuration,
+        event_manager=custom_event_manager,
+        storage_client=custom_storage_client,
+    )
+
+    assert service_locator.get_configuration() is custom_configuration
+    assert service_locator.get_event_manager() is custom_event_manager
+    assert service_locator.get_storage_client() is custom_storage_client
+
+    dataset = await crawler.get_dataset(name='test')
+    assert cast(DatasetClient, dataset._resource_client)._memory_storage_client is custom_storage_client
+
+
+async def test_allows_storage_client_overwrite_before_run(monkeypatch: pytest.MonkeyPatch) -> None:
+    custom_storage_client = MemoryStorageClient.from_config()
+
+    crawler = BasicCrawler(
+        storage_client=custom_storage_client,
+    )
+
+    @crawler.router.default_handler
+    async def handler(context: BasicCrawlingContext) -> None:
+        await context.push_data({'foo': 'bar'})
+
+    other_storage_client = MemoryStorageClient.from_config()
+    service_locator.set_storage_client(other_storage_client)
+
+    with monkeypatch.context() as monkey:
+        spy = Mock(wraps=service_locator.get_storage_client)
+        monkey.setattr(service_locator, 'get_storage_client', spy)
+        await crawler.run(['https://does-not-matter.com'])
+        assert spy.call_count >= 1
+
+    dataset = await crawler.get_dataset()
+    assert cast(DatasetClient, dataset._resource_client)._memory_storage_client is other_storage_client
+
+    data = await dataset.get_data()
+    assert data.items == [{'foo': 'bar'}]

tests/unit/conftest.py

@@ -42,9 +42,9 @@ def _prepare_test_env() -> None:
 
         # Reset the flags in the service locator to indicate that no services are explicitly set. This ensures
         # a clean state, as services might have been set during a previous test and not reset properly.
-        service_locator._configuration_was_set = False
-        service_locator._storage_client_was_set = False
-        service_locator._event_manager_was_set = False
+        service_locator._configuration_was_retrieved = False
+        service_locator._storage_client_was_retrieved = False
+        service_locator._event_manager_was_retrieved = False
 
         # Reset the services in the service locator.
         service_locator._configuration = None
@@ -61,9 +61,9 @@ def _prepare_test_env() -> None:
 
         # Verify that the test environment was set up correctly.
         assert os.environ.get('CRAWLEE_STORAGE_DIR') == str(tmp_path)
-        assert service_locator._configuration_was_set is False
-        assert service_locator._storage_client_was_set is False
-        assert service_locator._event_manager_was_set is False
+        assert service_locator._configuration_was_retrieved is False
+        assert service_locator._storage_client_was_retrieved is False
+        assert service_locator._event_manager_was_retrieved is False
 
     return _prepare_test_env
 

tests/unit/test_service_locator.py

@@ -9,41 +9,93 @@
 from crawlee.memory_storage_client import MemoryStorageClient
 
 
-def test_configuration() -> None:
+def test_default_configuration() -> None:
     default_config = Configuration()
     config = service_locator.get_configuration()
-    assert config == default_config
+    assert config == default_config  # == because these are in fact different instances, which should be fine
 
+
+def test_custom_configuration() -> None:
     custom_config = Configuration(default_browser_path='custom_path')
     service_locator.set_configuration(custom_config)
     config = service_locator.get_configuration()
-    assert config == custom_config
+    assert config is custom_config
+
+
+def test_configuration_overwrite() -> None:
+    default_config = Configuration()
+    service_locator.set_configuration(default_config)
+
+    custom_config = Configuration(default_browser_path='custom_path')
+    service_locator.set_configuration(custom_config)
+    assert service_locator.get_configuration() is custom_config
+
 
-    with pytest.raises(ServiceConflictError, match='Configuration has already been set.'):
+def test_configuration_conflict() -> None:
+    service_locator.get_configuration()
+    custom_config = Configuration(default_browser_path='custom_path')
+
+    with pytest.raises(ServiceConflictError, match='Configuration is already in use.'):
         service_locator.set_configuration(custom_config)
 
 
-def test_event_manager() -> None:
+def test_default_event_manager() -> None:
     default_event_manager = service_locator.get_event_manager()
     assert isinstance(default_event_manager, LocalEventManager)
 
+
+def test_custom_event_manager() -> None:
     custom_event_manager = LocalEventManager()
     service_locator.set_event_manager(custom_event_manager)
     event_manager = service_locator.get_event_manager()
-    assert event_manager == custom_event_manager
+    assert event_manager is custom_event_manager
+
+
+def test_event_manager_overwrite() -> None:
+    custom_event_manager = LocalEventManager()
+    service_locator.set_event_manager(custom_event_manager)
 
-    with pytest.raises(ServiceConflictError, match='EventManager has already been set.'):
+    another_custom_event_manager = LocalEventManager()
+    service_locator.set_event_manager(another_custom_event_manager)
+
+    assert custom_event_manager != another_custom_event_manager
+    assert service_locator.get_event_manager() is another_custom_event_manager
+
+
+def test_event_manager_conflict() -> None:
+    service_locator.get_event_manager()
+    custom_event_manager = LocalEventManager()
+
+    with pytest.raises(ServiceConflictError, match='EventManager is already in use.'):
         service_locator.set_event_manager(custom_event_manager)
 
 
-def test_storage_client() -> None:
+def test_default_storage_client() -> None:
     default_storage_client = service_locator.get_storage_client()
     assert isinstance(default_storage_client, MemoryStorageClient)
 
+
+def test_custom_storage_client() -> None:
     custom_storage_client = MemoryStorageClient.from_config()
     service_locator.set_storage_client(custom_storage_client)
     storage_client = service_locator.get_storage_client()
-    assert storage_client == custom_storage_client
+    assert storage_client is custom_storage_client
+
+
+def test_storage_client_overwrite() -> None:
+    custom_storage_client = MemoryStorageClient.from_config()
+    service_locator.set_storage_client(custom_storage_client)
+
+    another_custom_storage_client = MemoryStorageClient.from_config()
+    service_locator.set_storage_client(another_custom_storage_client)
+
+    assert custom_storage_client != another_custom_storage_client
+    assert service_locator.get_storage_client() is another_custom_storage_client
+
+
+def test_storage_client_conflict() -> None:
+    service_locator.get_storage_client()
+    custom_storage_client = MemoryStorageClient.from_config()
 
-    with pytest.raises(ServiceConflictError, match='StorageClient has already been set.'):
+    with pytest.raises(ServiceConflictError, match='StorageClient is already in use.'):
         service_locator.set_storage_client(custom_storage_client)