KubernetesJobWatcher no longer inherits from Process (#11017)

multiprocessing.Process is set up in a very unfortunate manner
that pretty much makes it impossible to test a class that inherits from
Process or use any of its internal functions. For this reason we decided
to seperate the actual process based functionality into a class member
This commit is contained in:
Daniel Imberman 2020-09-18 11:33:22 -07:00 коммит произвёл GitHub
Родитель 1a46e9b636
Коммит 1539bd051c
Не найден ключ, соответствующий данной подписи
Идентификатор ключа GPG: 4AEE18F83AFDEB23
2 изменённых файлов: 80 добавлений и 2 удалений

Просмотреть файл

@ -125,7 +125,7 @@ class KubeConfig: # pylint: disable=too-many-instance-attributes
return int(val)
class KubernetesJobWatcher(multiprocessing.Process, LoggingMixin):
class KubernetesJobWatcher(LoggingMixin):
"""Watches for Kubernetes jobs"""
def __init__(self,
@ -142,6 +142,31 @@ class KubernetesJobWatcher(multiprocessing.Process, LoggingMixin):
self.watcher_queue = watcher_queue
self.resource_version = resource_version
self.kube_config = kube_config
self.watcher_process = multiprocessing.Process(target=self.run, args=())
def start(self):
"""
Start the watcher process
"""
self.watcher_process.start()
def is_alive(self):
"""
Check if the watcher process is alive
"""
self.watcher_process.is_alive()
def join(self):
"""
Join watcher process
"""
self.watcher_process.join()
def terminate(self):
"""
Terminate watcher process
"""
self.watcher_process.terminate()
def run(self) -> None:
"""Performs watching"""

Просмотреть файл

@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
#
import multiprocessing
import random
import re
import string
@ -30,7 +31,9 @@ from tests.test_utils.config import conf_vars
try:
from kubernetes.client.rest import ApiException
from airflow.executors.kubernetes_executor import AirflowKubernetesScheduler, KubernetesExecutor
from airflow.executors.kubernetes_executor import (
AirflowKubernetesScheduler, KubernetesExecutor, KubernetesJobWatcher,
)
from airflow.kubernetes import pod_generator
from airflow.kubernetes.pod_generator import PodGenerator
from airflow.utils.state import State
@ -38,6 +41,56 @@ except ImportError:
AirflowKubernetesScheduler = None # type: ignore
class TestKubernetesJobWatcher(unittest.TestCase):
def setUp(self) -> None:
self.watcher_queue = multiprocessing.Manager().Queue()
self.watcher = KubernetesJobWatcher(
namespace="namespace",
multi_namespace_mode=False,
watcher_queue=self.watcher_queue,
resource_version="0",
worker_uuid="0",
kube_config=None,
)
def test_running_task(self):
self.watcher.process_status(
pod_id="pod_id",
namespace="namespace",
status="Running",
annotations={"foo": "bar"},
resource_version="5",
event={"type": "ADDED"}
)
self.assertTrue(self.watcher_queue.empty())
def test_succeeded_task(self):
self.watcher.process_status(
pod_id="pod_id",
namespace="namespace",
status="Succeeded",
annotations={"foo": "bar"},
resource_version="5",
event={"type": "ADDED"}
)
result = self.watcher_queue.get_nowait()
self.assertEqual(('pod_id', 'namespace', None, {'foo': 'bar'}, '5'), result)
self.assertTrue(self.watcher_queue.empty())
def test_failed_task(self):
self.watcher.process_status(
pod_id="pod_id",
namespace="namespace",
status="Failed",
annotations={"foo": "bar"},
resource_version="5",
event={"type": "ADDED"}
)
result = self.watcher_queue.get_nowait()
self.assertEqual(('pod_id', 'namespace', "failed", {'foo': 'bar'}, '5'), result)
self.assertTrue(self.watcher_queue.empty())
# pylint: disable=unused-argument
class TestAirflowKubernetesScheduler(unittest.TestCase):
@staticmethod