From b94810d044161ee4497a8e5eb23c873865657ce6 Mon Sep 17 00:00:00 2001
From: Ludeeus <ludeeus@ludeeus.dev>
Date: Fri, 27 Nov 2020 16:10:39 +0000
Subject: [PATCH] init

---
 supervisor/jobs/__init__.py    | 33 +++++++++++++++++++++++++++++++--
 supervisor/jobs/decorator.py   |  2 +-
 tests/jobs/test_job_context.py | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 3 deletions(-)
 create mode 100644 tests/jobs/test_job_context.py

diff --git a/supervisor/jobs/__init__.py b/supervisor/jobs/__init__.py
index 844511f75..c84b541e2 100644
--- a/supervisor/jobs/__init__.py
+++ b/supervisor/jobs/__init__.py
@@ -1,6 +1,8 @@
 """Supervisor job manager."""
+from contextvars import ContextVar
 import logging
 from typing import Dict, List, Optional
+from uuid import uuid4
 
 from ..coresys import CoreSys, CoreSysAttributes
 from ..utils.json import JsonConfig
@@ -8,18 +10,29 @@ from .const import ATTR_IGNORE_CONDITIONS, FILE_CONFIG_JOBS, JobCondition
 from .validate import SCHEMA_JOBS_CONFIG
 
 _LOGGER: logging.Logger = logging.getLogger(__package__)
+CONTEXT = ContextVar("id")
 
 
 class SupervisorJob(CoreSysAttributes):
     """Supervisor running job class."""
 
-    def __init__(self, coresys: CoreSys, name: str):
+    def __init__(self, coresys: CoreSys, name: Optional[str] = None):
         """Initialize the JobManager class."""
         self.coresys: CoreSys = coresys
-        self.name: str = name
+        self._name: Optional[str] = name
         self._progress: int = 0
         self._stage: Optional[str] = None
 
+    @property
+    def id(self) -> str:
+        """Return the ID for the job."""
+        return self.sys_jobs.context.get()
+
+    @property
+    def name(self) -> str:
+        """Return the name for the job."""
+        return self._name or self.id
+
     @property
     def progress(self) -> int:
         """Return the current progress."""
@@ -56,6 +69,7 @@ class JobManager(JsonConfig, CoreSysAttributes):
         """Initialize the JobManager class."""
         super().__init__(FILE_CONFIG_JOBS, SCHEMA_JOBS_CONFIG)
         self.coresys: CoreSys = coresys
+        self.context = CONTEXT
         self._jobs: Dict[str, SupervisorJob] = {}
 
     @property
@@ -68,6 +82,21 @@ class JobManager(JsonConfig, CoreSysAttributes):
         """Return a list of ingore condition."""
         return self._data[ATTR_IGNORE_CONDITIONS]
 
+    @property
+    def job(self) -> SupervisorJob:
+        """Return the current job based on context ID."""
+        try:
+            self.context.get()
+        except LookupError:
+            self.context.set(str(uuid4().hex))
+
+        context = self.context.get()
+
+        if context not in self._jobs:
+            self._jobs[context] = SupervisorJob(self.coresys)
+
+        return self._jobs[context]
+
     @ignore_conditions.setter
     def ignore_conditions(self, value: List[JobCondition]) -> None:
         """Set a list of ignored condition."""
diff --git a/supervisor/jobs/decorator.py b/supervisor/jobs/decorator.py
index 05b39dd78..29c70ba20 100644
--- a/supervisor/jobs/decorator.py
+++ b/supervisor/jobs/decorator.py
@@ -45,7 +45,7 @@ class Job:
             if not self._coresys:
                 raise JobException(f"coresys is missing on {self.name}")
 
-            job = self._coresys.jobs.get_job(self.name)
+            job = self._coresys.jobs.job
 
             if self.conditions and not self._check_conditions():
                 return False
diff --git a/tests/jobs/test_job_context.py b/tests/jobs/test_job_context.py
new file mode 100644
index 000000000..1dbec5867
--- /dev/null
+++ b/tests/jobs/test_job_context.py
@@ -0,0 +1,33 @@
+"""Test the job context vars."""
+
+from supervisor.coresys import CoreSys
+from supervisor.jobs.decorator import Job
+
+
+async def test_job_context(coresys: CoreSys):
+    """Test the job context."""
+
+    class TestClass:
+        """Test class."""
+
+        def __init__(self, coresys: CoreSys):
+            """Initialize the test class."""
+            self.coresys = coresys
+            self.jobid = str(self.coresys.jobs.job.id)
+
+        @Job()
+        async def first(self):
+            """Execute the class method."""
+            await self.second()
+            return True
+
+        @Job()
+        async def second(self):
+            """Execute the class method."""
+            return True
+
+    test = TestClass(coresys)
+    await test.first()
+    assert test.jobid == coresys.jobs.job.id
+    await test.second()
+    assert test.jobid == coresys.jobs.job.id