From 68f31db03ad03122213faa113588760c79819542 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Mon, 9 May 2022 16:23:35 -0400
Subject: [PATCH] UT: Add StubDIContainer

---
 monkey/common/__init__.py                     |  2 +-
 monkey/tests/common/__init__.py               |  1 +
 monkey/tests/common/stub_di_container.py      | 20 +++++++++++++++++++
 .../cc/resources/test_pba_file_download.py    |  4 ++--
 .../cc/resources/test_pba_file_upload.py      |  4 ++--
 5 files changed, 26 insertions(+), 5 deletions(-)
 create mode 100644 monkey/tests/common/__init__.py
 create mode 100644 monkey/tests/common/stub_di_container.py

diff --git a/monkey/common/__init__.py b/monkey/common/__init__.py
index 5f2cba2f1..01be39cf4 100644
--- a/monkey/common/__init__.py
+++ b/monkey/common/__init__.py
@@ -1 +1 @@
-from .di_container import DIContainer
+from .di_container import DIContainer, UnregisteredTypeError
diff --git a/monkey/tests/common/__init__.py b/monkey/tests/common/__init__.py
new file mode 100644
index 000000000..b50d9a82f
--- /dev/null
+++ b/monkey/tests/common/__init__.py
@@ -0,0 +1 @@
+from .stub_di_container import StubDIContainer
diff --git a/monkey/tests/common/stub_di_container.py b/monkey/tests/common/stub_di_container.py
new file mode 100644
index 000000000..bed434145
--- /dev/null
+++ b/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()
diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py
index d0a8ec48f..e05ab3362 100644
--- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py
+++ b/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:
diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py
index da1eb2c02..9cbaa50d8 100644
--- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py
+++ b/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: