Decouple scripts from WorkContext

tests/script/test_script.py

@@ -0,0 +1,133 @@
+from functools import partial
+import json
+import pytest
+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
+
+
+@pytest.mark.skipif(sys.version_info < (3, 8), reason="AsyncMock requires python 3.8+")
+class TestScript:
+    @pytest.fixture(autouse=True)
+    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}"
+
+    async def _on_download(self, expected, data: bytes):
+        assert data == expected
+        self._on_download_executed = True
+
+    @pytest.mark.asyncio
+    async def test_send_json(self, work_context: WorkContext):
+        storage: mock.AsyncMock = work_context._storage
+        dst_path = "/test/path"
+        data = {
+            "param": "value",
+        }
+
+        script = work_context.new_script()
+        script.send_json(data, dst_path)
+        await script._before()
+
+        storage.upload_bytes.assert_called_with(json.dumps(data).encode("utf-8"))
+        self._assert_dst_path(script, dst_path)
+
+    @pytest.mark.asyncio
+    async def test_send_bytes(self, work_context: WorkContext):
+        storage: mock.AsyncMock = work_context._storage
+        dst_path = "/test/path"
+        data = b"some byte string"
+
+        script = work_context.new_script()
+        script.send_bytes(data, dst_path)
+        await script._before()
+
+        storage.upload_bytes.assert_called_with(data)
+        self._assert_dst_path(script, dst_path)
+
+    @pytest.mark.asyncio
+    async def test_download_bytes(self, work_context: WorkContext):
+        expected = b"some byte string"
+        storage: mock.AsyncMock = work_context._storage
+        storage.new_destination.return_value.download_bytes.return_value = expected
+        src_path = "/test/path"
+
+        script = work_context.new_script()
+        script.download_bytes(src_path, partial(self._on_download, expected))
+        await script._before()
+        await script._after()
+
+        self._assert_src_path(script, src_path)
+        assert self._on_download_executed
+
+    @pytest.mark.asyncio
+    async def test_download_json(self, work_context: WorkContext):
+        expected = {"key": "val"}
+        storage: mock.AsyncMock = work_context._storage
+        storage.new_destination.return_value.download_bytes.return_value = json.dumps(
+            expected
+        ).encode("utf-8")
+        src_path = "/test/path"
+
+        script = work_context.new_script()
+        script.download_json(src_path, partial(self._on_download, expected))
+        await script._before()
+        await script._after()
+
+        self._assert_src_path(script, src_path)
+        assert self._on_download_executed
+
+    @pytest.mark.asyncio
+    async def test_implicit_init(self, work_context: WorkContext):
+        script = work_context.new_script()
+
+        # first script, should include implicit deploy and start cmds
+        await script._before()
+        assert len(script._commands) == 2
+        deploy_cmd, _ = script._commands[0]
+        assert isinstance(deploy_cmd, Deploy)
+        start_cmd, _ = script._commands[1]
+        assert isinstance(start_cmd, Start)
+        assert work_context._started
+
+        # second script, should not include implicit deploy and start
+        script = work_context.new_script()
+        script.run("/some/cmd")
+        await script._before()
+        assert len(script._commands) == 1
+
+    @pytest.mark.asyncio
+    async def test_cmd_result(self, work_context: WorkContext):
+        script = work_context.new_script()
+        future_result = script.run("/some/cmd", 1)
+
+        await script._before()
+        run_cmd, _ = script._commands[2]
+        result = CommandExecuted(
+            "job_id", "agr_id", "script_id", 2, command=run_cmd.evaluate(work_context)
+        )
+        script._set_cmd_result(result)
+
+        assert future_result.done()
+        assert future_result.result() == result

tests/test_ctx.py

@@ -6,6 +6,7 @@
 from unittest import mock
 
 from yapapi.ctx import CommandContainer, WorkContext
+from yapapi.script import Script
 
 
 def test_command_container():
@@ -37,6 +38,7 @@ def test_command_container():
     assert json.loads(expected_commands) == c.commands()
 
 
+@pytest.mark.skipif(sys.version_info < (3, 8), reason="AsyncMock requires python 3.8+")
 class TestWorkContext:
     @pytest.fixture(autouse=True)
     def setUp(self):
@@ -47,82 +49,85 @@ def _get_work_context(storage=None):
         return WorkContext(mock.Mock(), mock.Mock(), storage=storage)
 
     @staticmethod
-    def _assert_dst_path(steps, dst_path):
-        c = CommandContainer()
-        steps.register(c)
-        assert c.commands().pop()["transfer"]["to"] == f"container:{dst_path}"
+    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(steps, src_path):
-        c = CommandContainer()
-        steps.register(c)
-        assert c.commands().pop()["transfer"]["from"] == f"container:{src_path}"
+    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}"
 
     async def _on_download(self, expected, data: bytes):
         assert data == expected
         self._on_download_executed = True
 
     @pytest.mark.asyncio
-    @pytest.mark.skipif(sys.version_info < (3, 8), reason="AsyncMock requires python 3.8+")
     async def test_send_json(self):
         storage = mock.AsyncMock()
         dst_path = "/test/path"
         data = {
             "param": "value",
         }
         ctx = self._get_work_context(storage)
+
         ctx.send_json(dst_path, data)
-        steps = ctx.commit()
-        await steps.prepare()
+        script = ctx.commit()
+        await script._before()
+
         storage.upload_bytes.assert_called_with(json.dumps(data).encode("utf-8"))
-        self._assert_dst_path(steps, dst_path)
+        self._assert_dst_path(script, dst_path)
 
     @pytest.mark.asyncio
-    @pytest.mark.skipif(sys.version_info < (3, 8), reason="AsyncMock requires python 3.8+")
     async def test_send_bytes(self):
         storage = mock.AsyncMock()
         dst_path = "/test/path"
         data = b"some byte string"
         ctx = self._get_work_context(storage)
+
         ctx.send_bytes(dst_path, data)
-        steps = ctx.commit()
-        await steps.prepare()
+        script = ctx.commit()
+        await script._before()
+
         storage.upload_bytes.assert_called_with(data)
-        self._assert_dst_path(steps, dst_path)
+        self._assert_dst_path(script, dst_path)
 
     @pytest.mark.asyncio
-    @pytest.mark.skipif(sys.version_info < (3, 8), reason="AsyncMock requires python 3.8+")
     async def test_download_bytes(self):
         expected = b"some byte string"
-
         storage = mock.AsyncMock()
         storage.new_destination.return_value.download_bytes.return_value = expected
-
         src_path = "/test/path"
         ctx = self._get_work_context(storage)
+
         ctx.download_bytes(src_path, partial(self._on_download, expected))
-        steps = ctx.commit()
-        await steps.prepare()
-        await steps.post()
-        self._assert_src_path(steps, src_path)
+        script = ctx.commit()
+        await script._before()
+        await script._after()
+
+        self._assert_src_path(script, src_path)
         assert self._on_download_executed
 
     @pytest.mark.asyncio
-    @pytest.mark.skipif(sys.version_info < (3, 8), reason="AsyncMock requires python 3.8+")
     async def test_download_json(self):
         expected = {"key": "val"}
-
         storage = mock.AsyncMock()
         storage.new_destination.return_value.download_bytes.return_value = json.dumps(
             expected
         ).encode("utf-8")
         src_path = "/test/path"
         ctx = self._get_work_context(storage)
+
         ctx.download_json(src_path, partial(self._on_download, expected))
-        steps = ctx.commit()
-        await steps.prepare()
-        await steps.post()
-        self._assert_src_path(steps, src_path)
+        script = ctx.commit()
+        await script._before()
+        await script._after()
+
+        self._assert_src_path(script, src_path)
         assert self._on_download_executed
 
     @pytest.mark.parametrize(
@@ -135,19 +140,17 @@ async def test_download_json(self):
     def test_start(self, args):
         ctx = self._get_work_context()
         ctx.start(*args)
-        steps = ctx.commit()
+        script = ctx.commit()
 
-        c = CommandContainer()
-        steps.register(c)
+        batch = script._evaluate()
 
-        assert c.commands() == [{"start": {"args": args}}]
+        assert batch == [{"start": {"args": args}}]
 
     def test_terminate(self):
         ctx = self._get_work_context(None)
         ctx.terminate()
-        steps = ctx.commit()
+        script = ctx.commit()
 
-        c = CommandContainer()
-        steps.register(c)
+        batch = script._evaluate()
 
-        assert c.commands() == [{"terminate": {}}]
+        assert batch == [{"terminate": {}}]

yapapi/__init__.py

@@ -7,7 +7,7 @@
 from pkg_resources import get_distribution
 
 from yapapi.ctx import ExecOptions, Work, WorkContext
-from yapapi.engine import NoPaymentAccountError, WorkItem
+from yapapi.engine import NoPaymentAccountError
 from yapapi.executor import Executor, Task
 from yapapi.golem import Golem
 
@@ -55,5 +55,4 @@ class _WindowsEventPolicy(asyncio.events.BaseDefaultEventLoopPolicy):
     "Golem",
     "NoPaymentAccountError",
     "Work",
-    "WorkItem",
 ]