From 9c5283abfd04baaf04e27166905839bd544b6d5b Mon Sep 17 00:00:00 2001
From: Oded Messer <odedmesser@gmail.com>
Date: Sat, 13 Jul 2024 20:51:36 +0300
Subject: [PATCH] [Image] Docker run restart policy support (#40)

---
 .../runConfigurations/_it__basic_commands.xml |  16 +-
 .../runConfigurations/_it__shell_commands.xml |  14 +-
 .idea/runConfigurations/_ut__manof.xml        |  16 +-
 Makefile                                      |   2 +-
 manof/image.py                                |  10 +-
 pyproject.toml                                |  21 +-
 .../shell_commands/artifacts/manofest.py      |  21 +-
 .../shell_commands/test_shell_commands.py     | 181 +++++++++++-------
 8 files changed, 195 insertions(+), 86 deletions(-)

diff --git a/.idea/runConfigurations/_it__basic_commands.xml b/.idea/runConfigurations/_it__basic_commands.xml
index 7972fa2..795fb01 100644
--- a/.idea/runConfigurations/_it__basic_commands.xml
+++ b/.idea/runConfigurations/_it__basic_commands.xml
@@ -3,15 +3,25 @@
     <module name="manof" />
     <option name="INTERPRETER_OPTIONS" value="" />
     <option name="PARENT_ENVS" value="true" />
-    <option name="SDK_HOME" value="" />
+    <option name="SDK_HOME" value="$PROJECT_DIR$/venv/bin/python" />
     <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/tests/integration" />
-    <option name="IS_MODULE_SDK" value="true" />
+    <option name="IS_MODULE_SDK" value="false" />
     <option name="ADD_CONTENT_ROOTS" value="true" />
     <option name="ADD_SOURCE_ROOTS" value="true" />
     <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+    <EXTENSION ID="net.ashald.envfile">
+      <option name="IS_ENABLED" value="false" />
+      <option name="IS_SUBST" value="false" />
+      <option name="IS_PATH_MACRO_SUPPORTED" value="false" />
+      <option name="IS_IGNORE_MISSING_FILES" value="false" />
+      <option name="IS_ENABLE_EXPERIMENTAL_INTEGRATIONS" value="false" />
+      <ENTRIES>
+        <ENTRY IS_ENABLED="true" PARSER="runconfig" />
+      </ENTRIES>
+    </EXTENSION>
     <option name="_new_additionalArguments" value="&quot;&quot;" />
     <option name="_new_target" value="&quot;$PROJECT_DIR$/tests/integration/cases/basic/test_basic_commands.py&quot;" />
     <option name="_new_targetType" value="&quot;PATH&quot;" />
     <method v="2" />
   </configuration>
-</component>
\ No newline at end of file
+</component>
diff --git a/.idea/runConfigurations/_it__shell_commands.xml b/.idea/runConfigurations/_it__shell_commands.xml
index 990c42d..376c432 100644
--- a/.idea/runConfigurations/_it__shell_commands.xml
+++ b/.idea/runConfigurations/_it__shell_commands.xml
@@ -3,12 +3,22 @@
     <module name="manof" />
     <option name="INTERPRETER_OPTIONS" value="" />
     <option name="PARENT_ENVS" value="true" />
-    <option name="SDK_HOME" value="" />
+    <option name="SDK_HOME" value="$PROJECT_DIR$/venv/bin/python" />
     <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/tests/integration" />
-    <option name="IS_MODULE_SDK" value="true" />
+    <option name="IS_MODULE_SDK" value="false" />
     <option name="ADD_CONTENT_ROOTS" value="true" />
     <option name="ADD_SOURCE_ROOTS" value="true" />
     <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+    <EXTENSION ID="net.ashald.envfile">
+      <option name="IS_ENABLED" value="false" />
+      <option name="IS_SUBST" value="false" />
+      <option name="IS_PATH_MACRO_SUPPORTED" value="false" />
+      <option name="IS_IGNORE_MISSING_FILES" value="false" />
+      <option name="IS_ENABLE_EXPERIMENTAL_INTEGRATIONS" value="false" />
+      <ENTRIES>
+        <ENTRY IS_ENABLED="true" PARSER="runconfig" />
+      </ENTRIES>
+    </EXTENSION>
     <option name="_new_additionalArguments" value="&quot;&quot;" />
     <option name="_new_target" value="&quot;$PROJECT_DIR$/tests/integration/cases/shell_commands/test_shell_commands.py&quot;" />
     <option name="_new_targetType" value="&quot;PATH&quot;" />
diff --git a/.idea/runConfigurations/_ut__manof.xml b/.idea/runConfigurations/_ut__manof.xml
index bd29efb..0bdcc06 100644
--- a/.idea/runConfigurations/_ut__manof.xml
+++ b/.idea/runConfigurations/_ut__manof.xml
@@ -3,15 +3,25 @@
     <module name="manof" />
     <option name="INTERPRETER_OPTIONS" value="" />
     <option name="PARENT_ENVS" value="true" />
-    <option name="SDK_HOME" value="" />
+    <option name="SDK_HOME" value="$PROJECT_DIR$/venv/bin/python" />
     <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/tests/unit" />
-    <option name="IS_MODULE_SDK" value="true" />
+    <option name="IS_MODULE_SDK" value="false" />
     <option name="ADD_CONTENT_ROOTS" value="true" />
     <option name="ADD_SOURCE_ROOTS" value="true" />
     <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+    <EXTENSION ID="net.ashald.envfile">
+      <option name="IS_ENABLED" value="false" />
+      <option name="IS_SUBST" value="false" />
+      <option name="IS_PATH_MACRO_SUPPORTED" value="false" />
+      <option name="IS_IGNORE_MISSING_FILES" value="false" />
+      <option name="IS_ENABLE_EXPERIMENTAL_INTEGRATIONS" value="false" />
+      <ENTRIES>
+        <ENTRY IS_ENABLED="true" PARSER="runconfig" />
+      </ENTRIES>
+    </EXTENSION>
     <option name="_new_additionalArguments" value="&quot;&quot;" />
     <option name="_new_target" value="&quot;$PROJECT_DIR$/tests/unit/test_manof.py&quot;" />
     <option name="_new_targetType" value="&quot;PATH&quot;" />
     <method v="2" />
   </configuration>
-</component>
\ No newline at end of file
+</component>
diff --git a/Makefile b/Makefile
index da4fe6a..018182e 100644
--- a/Makefile
+++ b/Makefile
@@ -15,7 +15,7 @@ fmt:
 .PHONY: fmt-check
 fmt-check:
 	@echo "Running black fmt check..."
-	$(VENV_PYTHON) -m black --skip-string-normalization --check --diff -S --exclude='.*venv.*' .
+	$(VENV_PYTHON) -m black --skip-string-normalization --check --diff .
 
 .PHONY: test
 test: test-unit test-integ
diff --git a/manof/image.py b/manof/image.py
index b0505e0..4324d13 100644
--- a/manof/image.py
+++ b/manof/image.py
@@ -180,7 +180,7 @@ def run(self):
         for env in self._update_env_override():
 
             # single environment variable means set it to itself (x=x), thus forwarding the variable from the
-            # outter env into the docker env
+            # outer env into the docker env
             if isinstance(env, str):
                 lvalue = env
                 rvalue = os.environ.get(lvalue, None)
@@ -212,6 +212,10 @@ def run(self):
             if cap:
                 command += '--cap-drop={0} '.format(cap)
 
+        # set restart policy
+        if self.restart:
+            command += '--restart={0} '.format(self.restart)
+
         command = self._add_device_arguments(command)
 
         # set tag
@@ -727,6 +731,10 @@ def device_write_iops(self):
 
         return None
 
+    @property
+    def restart(self):
+        return None
+
     def to_dict(self):
         d = super(Image, self).to_dict()
         for idx, item in enumerate(d['volumes']):
diff --git a/pyproject.toml b/pyproject.toml
index 09c5ccf..c6358d6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,9 +1,22 @@
+[project]
+name = "manof"
+requires-python = ">=3.7,<3.10"
+version = "0.1.2"
+authors = [
+    {name = "LiranBG", email = "liran_ben_gida@mckinsey.com"},
+]
+
+[tool.setuptools]
+py-modules = []
+
 # please keep this to a minimum - defaults are good
 [tool.black]
 exclude = '''
-/(
+^/(   # anchor to the root
+  (   # exclude these directories
     \.git
-  | \.venv
-  | \venv
-)/
+    | \.venv
+    | venv
+  )/
+)
 '''
diff --git a/tests/integration/cases/shell_commands/artifacts/manofest.py b/tests/integration/cases/shell_commands/artifacts/manofest.py
index 0c30807..bf30ed4 100644
--- a/tests/integration/cases/shell_commands/artifacts/manofest.py
+++ b/tests/integration/cases/shell_commands/artifacts/manofest.py
@@ -26,7 +26,25 @@ def image_name(self):
 
     @property
     def command(self):
-        return '/bin/sh -c "echo \'{0}\'"'.format(self.name)
+        return '/bin/sh -c "echo \'{0}\' && sleep infinity"'.format(self.name)
+
+
+class TestImage2(TestImage):
+    @property
+    def restart(self):
+        return 'on-failure:5'
+
+    @property
+    def memory(self):
+        return '6Mib'
+
+    @property
+    def cpus(self):
+        return '1'
+
+    @property
+    def cap_add(self):
+        return ['SYS_ADMIN']
 
 
 class SomeGroup(manof.Group):
@@ -34,4 +52,5 @@ class SomeGroup(manof.Group):
     def members(self):
         return [
             'TestImage',
+            'TestImage2',
         ]
diff --git a/tests/integration/cases/shell_commands/test_shell_commands.py b/tests/integration/cases/shell_commands/test_shell_commands.py
index 6437c9f..06f0186 100644
--- a/tests/integration/cases/shell_commands/test_shell_commands.py
+++ b/tests/integration/cases/shell_commands/test_shell_commands.py
@@ -9,7 +9,7 @@
 class BasicCommandsTestCase(tests.integration.IntegrationTestCase):
     @defer.inlineCallbacks
     def test_serialize(self):
-        serialized_group_contents, _, _ = yield self._manof_command(
+        serialized_group_contents, _, _ = yield self._execute_manof_command(
             '--log-console-severity E serialize',
             [
                 'SomeGroup',
@@ -26,128 +26,167 @@ def test_run_verify_md5(self):
 
         # run twice to ensure md5 won't change between runs
         for _ in range(2):
-            yield self._manof_command('run', ['--dummy', 'do', target_name])
+            yield self._execute_manof_command('run', ['--dummy', 'do', target_name])
             command_sha = yield manof.utils.get_running_container_label(
                 target_name, label_name, self._logger
             )
             self.assertEqual('4a738101122b28baae05fac7a5dc6b32', command_sha)
 
-            run_md5, _, _ = yield self._manof_command(
+            run_md5, _, _ = yield self._execute_manof_command(
                 'run', ['--print-run-md5-only', '--dummy', 'do', target_name]
             )
             self.assertEqual('4a738101122b28baae05fac7a5dc6b32', run_md5)
 
         # run again and make ensure md5 has changed due to "--dummy" value change
-        yield self._manof_command('run', ['--dummy', 'else', target_name])
+        yield self._execute_manof_command('run', ['--dummy', 'value', target_name])
         command_sha = yield manof.utils.get_running_container_label(
             target_name, label_name, self._logger
         )
-        run_md5, _, _ = yield self._manof_command(
-            'run', ['--print-run-md5-only', '--dummy', 'else', target_name]
+        run_md5, _, _ = yield self._execute_manof_command(
+            'run', ['--print-run-md5-only', '--dummy', 'value', target_name]
         )
         self.assertEqual('a3ada1db9e167a8a747c8ddd4de63757', command_sha)
         self.assertEqual('a3ada1db9e167a8a747c8ddd4de63757', run_md5)
 
         # different dummy data yields different run md5
-        run_md5, _, _ = yield self._manof_command(
-            'run', ['--print-run-md5-only', '--dummy', 'else2', target_name]
+        run_md5, _, _ = yield self._execute_manof_command(
+            'run', ['--print-run-md5-only', '--dummy', 'value2', target_name]
         )
         self.assertNotEqual(run_md5, command_sha)
 
     @defer.inlineCallbacks
-    def test_run_and_rm(self):
-        self._logger.info('Testing run command happy flow')
+    def test_run_verify_md5(self):
+        self._logger.info('Testing run verify md5')
+        target_name = 'test_image'
+        label_name = manof.image.Constants.RUN_COMMAND_MD5_HASH_LABEL_NAME
 
-        image_name = 'test_image'
+        # run twice to ensure md5 won't change between runs
+        for _ in range(2):
+            yield self._execute_manof_command('run', ['--dummy', 'do', target_name])
+            command_sha = yield manof.utils.get_running_container_label(
+                target_name, label_name, self._logger
+            )
 
-        # sanity - removing containers if exists
-        yield self._remove_docker_container(image_name)
+            run_md5, _, _ = yield self._execute_manof_command(
+                'run', ['--print-run-md5-only', '--dummy', 'do', target_name]
+            )
+            self.assertEqual(command_sha, run_md5)
 
-        # run the image using manof
-        yield self._manof_command('run', [image_name])
+        # run again and make ensure md5 has changed due to "--dummy" value change
+        yield self._execute_manof_command('run', ['--dummy', 'else', target_name])
+        command_sha = yield manof.utils.get_running_container_label(
+            target_name, label_name, self._logger
+        )
+        run_md5, _, _ = yield self._execute_manof_command(
+            'run', ['--print-run-md5-only', '--dummy', 'else', target_name]
+        )
+        self.assertEqual(command_sha, run_md5)
 
-        # check the image exists
-        docker_log_output, _, _ = yield manof.utils.execute(
-            'docker logs {0}'.format(image_name),
-            cwd=None,
-            quiet=False,
-            logger=self._logger,
+        # different dummy data yields different run md5
+        run_md5, _, _ = yield self._execute_manof_command(
+            'run', ['--print-run-md5-only', '--dummy', 'else2', target_name]
         )
+        self.assertNotEqual(run_md5, command_sha)
+
+    @defer.inlineCallbacks
+    def test_image_run_and_rm(self):
+        self._logger.info('Testing run command happy flow')
+
+        for image_name in ['test_image', 'test_image2']:
 
-        self.assertEqual(docker_log_output, image_name)
+            # sanity - removing containers if exists
+            yield self._remove_docker_container(image_name)
 
-        # remove the container using manof
-        yield self._manof_command('rm', [image_name])
+            # run the image using manof
+            yield self._execute_manof_command('run', [image_name])
 
-        # check the container doesn't exist exists
-        yield self.assertFailure(
-            manof.utils.execute(
+            # check the container exists
+            docker_log_output, _, _ = yield manof.utils.execute(
                 'docker logs {0}'.format(image_name),
                 cwd=None,
                 quiet=False,
                 logger=self._logger,
-            ),
-            manof.utils.CommandFailedError,
-        )
-        self._logger.debug('Last command was supposed to fail')
+            )
+
+            self.assertEqual(docker_log_output, image_name)
+
+            # remove the container using manof
+            yield self._execute_manof_command('rm', [image_name, '--force'])
+
+            # check the container doesn't exist exists
+            yield self.assertFailure(
+                manof.utils.execute(
+                    'docker logs {0}'.format(image_name),
+                    cwd=None,
+                    quiet=False,
+                    logger=self._logger,
+                ),
+                manof.utils.CommandFailedError,
+            )
 
     @defer.inlineCallbacks
-    def test_provision_images(self):
+    def test_image_provision(self):
         self._logger.info('Testing provision images happy flow')
 
-        image_name = 'test_image'
-        manof_image = self._get_manof_image('TestImage')
-        docker_image = manof_image.image_name
+        for image_name, class_name in [
+            ('test_image', 'TestImage'),
+            ('test_image2', 'TestImage2'),
+        ]:
+            manof_image = self._get_manof_image(class_name)
+            docker_image = manof_image.image_name
 
-        # sanity - removing image if exists
-        yield self._remove_docker_image(docker_image)
+            # sanity - removing image if exists
+            yield self._remove_docker_image(docker_image)
 
-        # provision the image using manof
-        yield self._manof_command('provision', [image_name])
+            # provision the image using manof
+            yield self._execute_manof_command('provision', [image_name])
 
-        # check the image exists
-        yield manof.utils.execute(
-            'docker image history -Hq {0}'.format(docker_image),
-            cwd=None,
-            quiet=False,
-            logger=self._logger,
-        )
+            # check the image exists
+            yield manof.utils.execute(
+                'docker image history -Hq {0}'.format(docker_image),
+                cwd=None,
+                quiet=False,
+                logger=self._logger,
+            )
 
     @defer.inlineCallbacks
-    def test_lift_images(self):
+    def test_image_lift(self):
         self._logger.info('Testing provision images happy flow')
 
-        image_name = 'test_image'
-        manof_image = self._get_manof_image('TestImage')
-        docker_image = manof_image.image_name
+        for image_name, class_name in [
+            ('test_image', 'TestImage'),
+            ('test_image2', 'TestImage2'),
+        ]:
+            manof_image = self._get_manof_image(class_name)
+            docker_image = manof_image.image_name
 
-        # sanity - removing container and image if exists
-        yield self._remove_docker_container(image_name)
-        yield self._remove_docker_image(docker_image)
+            # sanity - removing container and image if exists
+            yield self._remove_docker_container(image_name)
+            yield self._remove_docker_image(docker_image)
 
-        # provision the image using manof
-        yield self._manof_command('lift', [image_name])
+            # provision the image using manof
+            yield self._execute_manof_command('lift', [image_name])
 
-        # check the image exists
-        yield manof.utils.execute(
-            'docker image history -Hq {0}'.format(docker_image),
-            cwd=None,
-            quiet=False,
-            logger=self._logger,
-        )
+            # check the image exists
+            yield manof.utils.execute(
+                'docker image history -Hq {0}'.format(docker_image),
+                cwd=None,
+                quiet=False,
+                logger=self._logger,
+            )
 
-        # check the image exists
-        docker_log_output, _, _ = yield manof.utils.execute(
-            'docker logs {0}'.format(image_name),
-            cwd=None,
-            quiet=False,
-            logger=self._logger,
-        )
+            # check the container exists
+            docker_log_output, _, _ = yield manof.utils.execute(
+                'docker logs {0}'.format(image_name),
+                cwd=None,
+                quiet=False,
+                logger=self._logger,
+            )
 
-        self.assertEqual(docker_log_output, image_name)
+            self.assertEqual(docker_log_output, image_name)
 
     @defer.inlineCallbacks
-    def _manof_command(self, command, args):
+    def _execute_manof_command(self, command, args):
         out, err, signal = yield manof.utils.execute(
             'manof {command} {args}'.format(command=command, args=' '.join(args)),
             cwd=self._working_dir,