From 2267f1af060a8d50aa1b60b566b63463ef7e1131 Mon Sep 17 00:00:00 2001
From: Till Riedel <riedel@teco.edu>
Date: Sun, 31 Oct 2021 16:49:29 +0100
Subject: [PATCH] add functions for testing within cluster ci add test for
 security inside ci

---
 conftest.py                               |  8 ++-
 dask_jobqueue/tests/test_jobqueue_core.py | 71 ++++++++++++++++++-----
 2 files changed, 64 insertions(+), 15 deletions(-)

diff --git a/conftest.py b/conftest.py
index a191a990..d6b27cf1 100644
--- a/conftest.py
+++ b/conftest.py
@@ -25,7 +25,13 @@ def pytest_configure(config):
 
 
 def pytest_runtest_setup(item):
-    envnames = [mark.args[0] for mark in item.iter_markers(name="env")]
+    envnames = sum(
+        [
+            mark.args[0] if isinstance(mark.args[0], list) else [mark.args[0]]
+            for mark in item.iter_markers(name="env")
+        ],
+        [],
+    )
     if (item.config.getoption("-E") is None and envnames) or (
         item.config.getoption("-E") is not None
         and item.config.getoption("-E") not in envnames
diff --git a/dask_jobqueue/tests/test_jobqueue_core.py b/dask_jobqueue/tests/test_jobqueue_core.py
index b616392c..40a76b17 100644
--- a/dask_jobqueue/tests/test_jobqueue_core.py
+++ b/dask_jobqueue/tests/test_jobqueue_core.py
@@ -27,15 +27,26 @@
 
 from dask_jobqueue.sge import SGEJob
 
-all_clusters = [
-    PBSCluster,
-    MoabCluster,
-    SLURMCluster,
-    SGECluster,
-    LSFCluster,
-    OARCluster,
-    HTCondorCluster,
-]
+all_envs = {
+    "pbs": PBSCluster,
+    "moab": MoabCluster,
+    "slurm": SLURMCluster,
+    "sge": SGECluster,
+    "lsf": LSFCluster,
+    "oar": OARCluster,
+    "htcondor": HTCondorCluster,
+}
+
+
+@pytest.fixture
+def EnvCluster(pytestconfig):
+    if pytestconfig.getoption("-E") is not None:
+        return all_envs[pytestconfig.getoption("-E")]
+    else:
+        return LocalCluster
+
+
+all_clusters = list(all_envs.values())
 
 
 def create_cluster_func(cluster_cls, **kwargs):
@@ -426,10 +437,7 @@ def test_security():
         tls_client_cert=cert,
         require_encryption=True,
     )
-
-    with LocalCluster(
-        cores=1, memory="1GB", security=security, protocol="tls"
-    ) as cluster:
+    with LocalCluster(cores=1, memory="1GB", security=security) as cluster:
         assert cluster.security == security
         assert cluster.scheduler_spec["options"]["security"] == security
         job_script = cluster.job_script()
@@ -444,5 +452,40 @@ def test_security():
             result = future.result()
             assert result == 11
 
-    with LocalCluster(cores=1, memory="1GB", security=security) as cluster:
+
+@pytest.mark.env(["htcondor", "slurm", "pbs", "sge"])
+def test_security_cluster(EnvCluster):
+    dirname = os.path.dirname(__file__)
+    key = os.path.join(dirname, "key.pem")
+    cert = os.path.join(dirname, "ca.pem")
+    security = Security(
+        tls_ca_file=cert,
+        tls_scheduler_key=key,
+        tls_scheduler_cert=cert,
+        tls_worker_key=key,
+        tls_worker_cert=cert,
+        tls_client_key=key,
+        tls_client_cert=cert,
+        require_encryption=True,
+    )
+
+    with create_cluster_func(
+        EnvCluster,
+        cores=1,
+        memory="100MB",
+        security=security,
+        protocol="tls",
+    ) as cluster:
+        assert cluster.security == security
+        assert cluster.scheduler_spec["options"]["security"] == security
+        job_script = cluster.job_script()
         assert "tls://" in job_script
+        assert "--tls-key {}".format(key) in job_script
+        assert "--tls-cert {}".format(cert) in job_script
+        assert "--tls-ca-file {}".format(cert) in job_script
+
+        cluster.scale(jobs=1)
+        with Client(cluster, security=security) as client:
+            future = client.submit(lambda x: x + 1, 10)
+            result = future.result()
+            assert result == 11