adds tests and small fixes

Author: fabricebrito

calrissian/context.py

@@ -23,6 +23,6 @@ def __init__(self, kwargs=None):
         self.pod_serviceaccount = None
         self.tool_logs_basepath = None
         self.max_gpus = None
-        self.no_network_access_pod_label = None
-        self.network_access_pod_label = None
+        self.no_network_access_pod_labels = None
+        self.network_access_pod_labels = None
         return super(CalrissianRuntimeContext, self).__init__(kwargs)

calrissian/job.py

@@ -198,7 +198,7 @@ def add_emptydir_volume_binding(self, name, target):
 
 class KubernetesPodBuilder(object):
 
-    def __init__(self, name, builder, container_image, environment, volume_mounts, volumes, command_line, stdout, stderr, stdin, labels, nodeselectors, security_context, serviceaccount, no_network_access_pod_label={}, network_access_pod_label={}):
+    def __init__(self, name, builder, container_image, environment, volume_mounts, volumes, command_line, stdout, stderr, stdin, labels, nodeselectors, security_context, serviceaccount, no_network_access_pod_labels=None, network_access_pod_labels=None):
         self.name = name
         self.builder = builder
         self.cwl_version = self.builder.cwlVersion
@@ -215,8 +215,8 @@ def __init__(self, name, builder, container_image, environment, volume_mounts, v
         self.nodeselectors = nodeselectors
         self.security_context = security_context
         self.serviceaccount = serviceaccount
-        self.no_network_access_pod_label = no_network_access_pod_label
-        self.network_access_pod_label = network_access_pod_label
+        self.no_network_access_pod_labels = no_network_access_pod_labels
+        self.network_access_pod_labels = network_access_pod_labels
         self.requirements = {} if self.builder.requirements is None else self.builder.requirements
         self.hints = [] if self.builder.hints is None else self.builder.hints
 
@@ -350,15 +350,14 @@ def pod_labels(self):
             network_access = False
 
         for requirement in self.requirements:
-            if requirement["class"] in ["NetworkAccess"]:
-                network_access = requirement.get("networkAccess")
+            if "class" in requirement.keys() and requirement["class"] in ["NetworkAccess"]:
+                network_access = True if requirement.get("networkAccess") == "true" else False
                 break
-        
-        if not network_access and self.no_network_access_pod_label: 
-            self.labels = {**self.labels, **self.no_network_access_pod_label}
+        if not network_access and self.no_network_access_pod_labels: 
+            self.labels = {**self.labels, **self.no_network_access_pod_labels}
 
-        if network_access and self.network_access_pod_label:
-            self.labels = {**self.labels, **self.network_access_pod_label}
+        if network_access and self.network_access_pod_labels:
+            self.labels = {**self.labels, **self.network_access_pod_labels}
 
         return {str(k): str(v) for k, v in self.labels.items()}
 
@@ -536,15 +535,15 @@ def get_pod_labels(self, runtimeContext):
         else:
             return {}
 
-    def get_network_access_pod_label(self, runtimeContext):
-        if runtimeContext.network_access_pod_label:
-            return read_yaml(runtimeContext.network_access_pod_label)
+    def get_network_access_pod_labels(self, runtimeContext):
+        if runtimeContext.network_access_pod_labels:
+            return read_yaml(runtimeContext.network_access_pod_labels)
         else:
             return {}
 
-    def get_no_network_access_pod_label(self, runtimeContext):
-        if runtimeContext.no_network_access_pod_label:
-            return read_yaml(runtimeContext.no_network_access_pod_label)
+    def get_no_network_access_pod_labels(self, runtimeContext):
+        if runtimeContext.no_network_access_pod_labels:
+            return read_yaml(runtimeContext.no_network_access_pod_labels)
         else:
             return {}
 
@@ -621,8 +620,8 @@ def create_kubernetes_runtime(self, runtimeContext):
             self.get_pod_nodeselectors(runtimeContext),
             self.get_security_context(runtimeContext),
             self.get_pod_serviceaccount(runtimeContext),
-            self.get_no_network_access_pod_label(runtimeContext),
-            self.get_network_access_pod_label(runtimeContext),
+            self.get_no_network_access_pod_labels(runtimeContext),
+            self.get_network_access_pod_labels(runtimeContext),
         )
         built = k8s_builder.build()
         log.debug('{}\n{}{}\n'.format('-' * 80, yaml.dump(built), '-' * 80))

tests/test_job.py

@@ -274,7 +274,7 @@ class KubernetesPodBuilderTestCase(TestCase):
 
     def setUp(self):
         builder = Mock()
-        builder.cwlVersion = "v1.0"
+        builder.cwlVersion = "v1.2"
         builder.requirements = []
         builder.resources = {'cores': 1, 'ram': 1024}
         self.name = 'PodName'
@@ -291,12 +291,12 @@ def setUp(self):
         self.nodeselectors = {'disktype': 'ssd', 'cachelevel': 2}
         self.security_context = { 'runAsUser': os.getuid(),'runAsGroup': os.getgid() }
         self.pod_serviceaccount = "podmanager"
-        self.no_network_access_pod_label = {}
-        self.network_access_pod_label = {}
+        self.no_network_access_pod_labels = {"calrissian-network": "disabled"}
+        self.network_access_pod_labels = {"calrissian-network": "enabled"}
         self.pod_builder = KubernetesPodBuilder(self.name, self.builder, self.container_image, self.environment, self.volume_mounts,
                                                 self.volumes, self.command_line, self.stdout, self.stderr, self.stdin,
                                                 self.labels, self.nodeselectors, self.security_context, self.pod_serviceaccount, 
-                                                self.no_network_access_pod_label, self.network_access_pod_label)
+                                                self.no_network_access_pod_labels, self.network_access_pod_labels)
 
     @patch('calrissian.job.random_tag')
     def test_safe_pod_name(self, mock_random_tag):
@@ -382,14 +382,37 @@ def test_gpu_hints(self):
         }
         self.assertEqual(expected, resources)
 
+    def test_network_access_1_2(self):
+        self.pod_builder.cwl_version = "v1.2"
+        self.pod_builder.requirements = [OrderedDict([("class", "NetworkAccess"), ("networkAccess", "true")])]
+        self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "enabled", 'key1':'val1', 'key2':'123'})
+
+    def test_no_network_access_1_2(self):
+        self.pod_builder.cwl_version = "v1.2"
+        self.pod_builder.requirements = [OrderedDict([("class", "NetworkAccess"), ("networkAccess", "false")])]
+        self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "disabled", 'key1':'val1', 'key2':'123'})
+
+    def test_network_access_1_0(self):
+        self.pod_builder.cwl_version = "v1.0"
+        self.pod_builder.requirements = [OrderedDict([])]
+        self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "enabled", 'key1':'val1', 'key2':'123'})
+
     def test_string_labels(self):
         self.pod_builder.labels = {'key1': 123}
-        self.assertEqual(self.pod_builder.pod_labels(), {'key1':'123'})
+        self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "disabled", 'key1':'123'})
 
     def test_string_nodeselectors(self):
         self.pod_builder.nodeselectors = {'cachelevel': 2}
         self.assertEqual(self.pod_builder.pod_nodeselectors(), {'cachelevel':'2'})
 
+    def test_string_no_network_access_pod_label(self):
+        self.pod_builder.no_network_access_pod_labels = {"calrissian-network": "disabled"}
+        self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "disabled", 'key1': 'val1', 'key2': '123'})
+
+    def test_string_network_access_pod_label(self):
+        self.pod_builder.network_access_pod_labels = {"calrissian-network": "enabled"}
+        self.assertEqual(self.pod_builder.pod_labels(), {"calrissian-network": "disabled", 'key1': 'val1', 'key2': '123'})
+
     def test_init_containers_empty_when_no_stdout_or_stderr(self):
         self.pod_builder.stdout = None
         self.pod_builder.stderr = None
@@ -435,6 +458,7 @@ def test_build(self, mock_random_tag):
                 'labels': {
                     'key1': 'val1',
                     'key2': '123',
+                    'calrissian-network': 'disabled',
                 }
             },
             'apiVersion': 'v1',