UT: Add StubDIContainer

monkey/common/__init__.py

@@ -1 +1 @@
-from .di_container import DIContainer
+from .di_container import DIContainer, UnregisteredTypeError

monkey/tests/common/__init__.py

@@ -0,0 +1 @@
+from .stub_di_container import StubDIContainer

monkey/tests/common/stub_di_container.py

@@ -0,0 +1,20 @@
+from typing import Any, Sequence, Type, TypeVar
+from unittest.mock import MagicMock
+
+from common import DIContainer, UnregisteredTypeError
+
+T = TypeVar("T")
+
+
+class StubDIContainer(DIContainer):
+    def resolve(self, type_: Type[T]) -> T:
+        try:
+            return super().resolve(type_)
+        except UnregisteredTypeError:
+            return MagicMock()
+
+    def resolve_dependencies(self, type_: Type[T]) -> Sequence[Any]:
+        try:
+            return super().resolve_dependencies(type_)
+        except UnregisteredTypeError:
+            return MagicMock()

monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py

@@ -2,8 +2,8 @@ import io
 from typing import BinaryIO
 
 import pytest
+from tests.common import StubDIContainer
 
-from common import DIContainer
 from monkey_island.cc.services import FileRetrievalError, IFileStorageService
 
 FILE_NAME = "test_file"
@@ -32,7 +32,7 @@ class MockFileStorageService(IFileStorageService):
 
 @pytest.fixture
 def flask_client(build_flask_client, tmp_path):
-    container = DIContainer()
+    container = StubDIContainer()
     container.register(IFileStorageService, MockFileStorageService)
 
     with build_flask_client(container) as flask_client:

monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py

@@ -2,9 +2,9 @@ import io
 from typing import BinaryIO
 
 import pytest
+from tests.common import StubDIContainer
 from tests.utils import raise_
 
-from common import DIContainer
 from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE
 from monkey_island.cc.services import FileRetrievalError, IFileStorageService
 
@@ -65,7 +65,7 @@ def file_storage_service():
 
 @pytest.fixture
 def flask_client(build_flask_client, file_storage_service):
-    container = DIContainer()
+    container = StubDIContainer()
     container.register_instance(IFileStorageService, file_storage_service)
 
     with build_flask_client(container) as flask_client: