From 53d23eb5994ec9be654de9a25b15b343a7a374f0 Mon Sep 17 00:00:00 2001
From: Kuba Mazurek <zakaprov@gmail.com>
Date: Wed, 8 Sep 2021 19:48:16 +0200
Subject: [PATCH] Add a factory for WorkContext objects in tests

---
 tests/factories/context.py  | 13 +++++++++++++
 tests/script/test_script.py | 28 +++++++++++++++-------------
 tests/test_ctx.py           |  2 --
 3 files changed, 28 insertions(+), 15 deletions(-)
 create mode 100644 tests/factories/context.py

diff --git a/tests/factories/context.py b/tests/factories/context.py
new file mode 100644
index 000000000..7257e79bb
--- /dev/null
+++ b/tests/factories/context.py
@@ -0,0 +1,13 @@
+import factory
+from unittest import mock
+
+from yapapi.ctx import WorkContext
+
+
+class WorkContextFactory(factory.Factory):
+    class Meta:
+        model = WorkContext
+
+    activity = mock.MagicMock()
+    agreement_details = mock.MagicMock()
+    storage = mock.AsyncMock()
diff --git a/tests/script/test_script.py b/tests/script/test_script.py
index fcd1658db..d9d7e2bbf 100644
--- a/tests/script/test_script.py
+++ b/tests/script/test_script.py
@@ -4,11 +4,13 @@
 import sys
 from unittest import mock
 
-from yapapi import WorkContext
 from yapapi.events import CommandExecuted
 from yapapi.script import Script
 from yapapi.script.command import Deploy, Start
 
+if sys.version_info >= (3, 8):
+    from tests.factories.context import WorkContextFactory
+
 
 @pytest.mark.skipif(sys.version_info < (3, 8), reason="AsyncMock requires python 3.8+")
 class TestScript:
@@ -16,21 +18,15 @@ class TestScript:
     def setUp(self):
         self._on_download_executed = False
 
-    @pytest.fixture
-    def work_context(self):
-        return WorkContext(mock.MagicMock(), mock.MagicMock(), storage=mock.AsyncMock())
-
     @staticmethod
     def _assert_dst_path(script: Script, dst_path):
         batch = script._evaluate()
-        # transfer_cmd = {'transfer': {'from': 'some/mock/path', 'to': 'container:expected/path'}
         transfer_cmd = [cmd for cmd in batch if "transfer" in cmd][0]
         assert transfer_cmd["transfer"]["to"] == f"container:{dst_path}"
 
     @staticmethod
     def _assert_src_path(script: Script, src_path):
         batch = script._evaluate()
-        # transfer_cmd = {'transfer': {'from': 'container:expected/path', 'to': 'some/mock/path'}
         transfer_cmd = [cmd for cmd in batch if "transfer" in cmd][0]
         assert transfer_cmd["transfer"]["from"] == f"container:{src_path}"
 
@@ -39,7 +35,8 @@ async def _on_download(self, expected, data: bytes):
         self._on_download_executed = True
 
     @pytest.mark.asyncio
-    async def test_send_json(self, work_context: WorkContext):
+    async def test_send_json(self):
+        work_context = WorkContextFactory()
         storage: mock.AsyncMock = work_context._storage
         dst_path = "/test/path"
         data = {
@@ -54,7 +51,8 @@ async def test_send_json(self, work_context: WorkContext):
         self._assert_dst_path(script, dst_path)
 
     @pytest.mark.asyncio
-    async def test_send_bytes(self, work_context: WorkContext):
+    async def test_send_bytes(self):
+        work_context = WorkContextFactory()
         storage: mock.AsyncMock = work_context._storage
         dst_path = "/test/path"
         data = b"some byte string"
@@ -67,7 +65,8 @@ async def test_send_bytes(self, work_context: WorkContext):
         self._assert_dst_path(script, dst_path)
 
     @pytest.mark.asyncio
-    async def test_download_bytes(self, work_context: WorkContext):
+    async def test_download_bytes(self):
+        work_context = WorkContextFactory()
         expected = b"some byte string"
         storage: mock.AsyncMock = work_context._storage
         storage.new_destination.return_value.download_bytes.return_value = expected
@@ -82,7 +81,8 @@ async def test_download_bytes(self, work_context: WorkContext):
         assert self._on_download_executed
 
     @pytest.mark.asyncio
-    async def test_download_json(self, work_context: WorkContext):
+    async def test_download_json(self):
+        work_context = WorkContextFactory()
         expected = {"key": "val"}
         storage: mock.AsyncMock = work_context._storage
         storage.new_destination.return_value.download_bytes.return_value = json.dumps(
@@ -99,7 +99,8 @@ async def test_download_json(self, work_context: WorkContext):
         assert self._on_download_executed
 
     @pytest.mark.asyncio
-    async def test_implicit_init(self, work_context: WorkContext):
+    async def test_implicit_init(self):
+        work_context = WorkContextFactory()
         script = work_context.new_script()
 
         # first script, should include implicit deploy and start cmds
@@ -118,7 +119,8 @@ async def test_implicit_init(self, work_context: WorkContext):
         assert len(script._commands) == 1
 
     @pytest.mark.asyncio
-    async def test_cmd_result(self, work_context: WorkContext):
+    async def test_cmd_result(self):
+        work_context = WorkContextFactory()
         script = work_context.new_script()
         future_result = script.run("/some/cmd", 1)
 
diff --git a/tests/test_ctx.py b/tests/test_ctx.py
index becb2a795..80d8f8d12 100644
--- a/tests/test_ctx.py
+++ b/tests/test_ctx.py
@@ -51,14 +51,12 @@ def _get_work_context(storage=None):
     @staticmethod
     def _assert_dst_path(script: Script, dst_path):
         batch = script._evaluate()
-        # transfer_cmd = {'transfer': {'from': 'some/mock/path', 'to': 'container:expected/path'}
         transfer_cmd = [cmd for cmd in batch if "transfer" in cmd][0]
         assert transfer_cmd["transfer"]["to"] == f"container:{dst_path}"
 
     @staticmethod
     def _assert_src_path(script: Script, src_path):
         batch = script._evaluate()
-        # transfer_cmd = {'transfer': {'from': 'container:expected/path', 'to': 'some/mock/path'}
         transfer_cmd = [cmd for cmd in batch if "transfer" in cmd][0]
         assert transfer_cmd["transfer"]["from"] == f"container:{src_path}"