From a33c94e24f436faba89c6abd8c99b528e4c78cff Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Tue, 9 Jul 2019 00:04:26 +0300
Subject: [PATCH] Add task diff support

---
 server/database/model/task/task.py        |  2 ++
 server/tests/automated/test_tasks_diff.py | 38 +++++++++++++++++++++++
 server/utilities/threads_manager.py       | 29 +++++++++++++++++
 3 files changed, 69 insertions(+)
 create mode 100644 server/tests/automated/test_tasks_diff.py
 create mode 100644 server/utilities/threads_manager.py

diff --git a/server/database/model/task/task.py b/server/database/model/task/task.py
index 7db50d0..fbb788e 100644
--- a/server/database/model/task/task.py
+++ b/server/database/model/task/task.py
@@ -29,6 +29,7 @@ class TaskStatus(object):
     published = 'published'
     closed = 'closed'
     failed = 'failed'
+    completed = 'completed'
     unknown = 'unknown'
 
 
@@ -49,6 +50,7 @@ class Script(EmbeddedDocument):
     entry_point = StringField(required=True)
     working_dir = StringField()
     requirements = SafeDictField()
+    diff = StringField()
 
 
 class Execution(EmbeddedDocument):
diff --git a/server/tests/automated/test_tasks_diff.py b/server/tests/automated/test_tasks_diff.py
new file mode 100644
index 0000000..a059d98
--- /dev/null
+++ b/server/tests/automated/test_tasks_diff.py
@@ -0,0 +1,38 @@
+from config import config
+from tests.automated import TestService
+
+log = config.logger(__file__)
+
+
+class TestTasksDiff(TestService):
+
+    def setUp(self, version="2.0"):
+        super(TestTasksDiff, self).setUp(version=version)
+
+    def new_task(self, **kwargs):
+        return self.create_temp(
+            "tasks", name="test", type="testing", input=dict(view=dict()), **kwargs
+        )
+
+    def _compare_script(self, task, script):
+        for key, value in script.items():
+            self.assertEqual(task.script[key], value)
+
+    def test_not_deleted(self):
+        task_id = self.new_task()
+        script = dict(
+            requirements=dict(pip=["six"]),
+            repository="https://example.come/foo/bar",
+            entry_point="test.py",
+            diff="foo",
+        )
+        self.api.tasks.edit(task=task_id, script=script)
+        self.api.tasks.started(task=task_id)
+        self.api.tasks.reset(task=task_id)
+        task = self.api.tasks.get_by_id(task=task_id).task
+        self._compare_script(task, script)
+        new_reqs = dict()
+        self.api.tasks.set_requirements(task=task_id, requirements=new_reqs)
+        script["requirements"] = new_reqs
+        task = self.api.tasks.get_by_id(task=task_id).task
+        self._compare_script(task, script)
diff --git a/server/utilities/threads_manager.py b/server/utilities/threads_manager.py
new file mode 100644
index 0000000..f7ab2b5
--- /dev/null
+++ b/server/utilities/threads_manager.py
@@ -0,0 +1,29 @@
+from functools import wraps
+from threading import Lock, Thread
+
+import attr
+
+
+@attr.s(auto_attribs=True)
+class ThreadsManager:
+    objects = {}
+    lock = Lock()
+
+    def register(self, thread_name, daemon=True):
+        def decorator(f):
+            @wraps(f)
+            def wrapper(*args, **kwargs):
+                with self.lock:
+                    thread = self.objects.get(thread_name)
+                    if not thread:
+                        thread = Thread(
+                            target=f, name=thread_name, args=args, kwargs=kwargs
+                        )
+                        thread.daemon = daemon
+                        thread.start()
+                        self.objects[thread_name] = thread
+                    return thread.ident
+
+            return wrapper
+
+        return decorator