Fix recursively transforming collections to JSON (#308)

- and add testing for it

Author: tazend

pyslurm/core/job/job.pyx

@@ -295,7 +295,9 @@ cdef class Job:
         Returns:
             (dict): Job information as dict
         """
-        return instance_to_dict(self)
+        cdef dict out = instance_to_dict(self)
+        out["steps"] = self.steps.to_dict()
+        return out
 
     def send_signal(self, signal, steps="children", hurry=False):
         """Send a signal to a running Job.

pyslurm/core/job/step.pyx

@@ -143,6 +143,9 @@ cdef class JobSteps(dict):
         cdef JobSteps steps = JobSteps()
         return steps._load_data(slurm.NO_VAL, slurm.SHOW_ALL)
 
+    def to_dict(self):
+        return xcollections.dict_recursive(self)
+
 
 cdef class JobStep:
 
@@ -329,7 +332,13 @@ cdef class JobStep:
         Returns:
             (dict): JobStep information as dict
         """
-        return instance_to_dict(self)
+        cdef dict out = instance_to_dict(self)
+
+        dist = self.distribution
+        if dist:
+            out["distribution"] = dist.to_dict()
+
+        return out
 
     @property
     def id(self):

pyslurm/db/job.pyx

@@ -524,11 +524,8 @@ cdef class Job:
 
         if self.stats:
             out["stats"] = self.stats.to_dict()
-
-        steps = out.pop("steps", {})
-        out["steps"] = {}
-        for step_id, step in steps.items():
-            out["steps"][step_id] = step.to_dict() 
+        if self.steps:
+            out["steps"] = self.steps.to_dict()
 
         return out
 

pyslurm/db/step.pyx

@@ -27,6 +27,7 @@ from pyslurm.core.error import RPCError
 from typing import Union
 from pyslurm.utils.uint import *
 from pyslurm.utils.ctime import _raw_time
+from pyslurm import xcollections
 from pyslurm.utils.helpers import (
     gid_to_name,
     uid_to_name,
@@ -43,6 +44,9 @@ cdef class JobSteps(dict):
         data = super().__repr__()
         return f'pyslurm.db.{self.__class__.__name__}({data})'
 
+    def to_dict(self):
+        return xcollections.dict_recursive(self)
+
 
 cdef class JobStep:
 

pyslurm/xcollections.pyx

@@ -447,6 +447,9 @@ cdef class MultiClusterMap:
         Returns:
             (str): JSON formatted string from `json.dumps()`
         """
+        if not self.data:
+            return '{}'
+
         data = multi_dict_recursive(self)
         if multi_cluster:
             return json.dumps(data)

tests/integration/test_db_job.py

@@ -24,6 +24,7 @@
 import pyslurm
 import time
 import util
+import json
 
 
 # TODO: Instead of submitting new Jobs and waiting to test Database API
@@ -55,6 +56,20 @@ def test_parse_all(submit_job):
     assert job_dict["steps"]
 
 
+def test_to_json(submit_job):
+    job = submit_job()
+    util.wait()
+
+    jfilter = pyslurm.db.JobFilter(ids=[job.id])
+    jobs = pyslurm.db.Jobs.load(jfilter)
+
+    json_data = jobs.to_json()
+    dict_data = json.loads(json_data)
+    assert dict_data
+    assert json_data
+    assert len(dict_data) == 1
+
+
 def test_modify(submit_job):
     job = submit_job()
     util.wait(5)

tests/integration/test_job.py

@@ -23,6 +23,7 @@
 import time
 import pytest
 import pyslurm
+import json
 import util
 from util import create_simple_job_desc
 from pyslurm import (
@@ -64,34 +65,34 @@ def test_cancel(submit_job):
     job = submit_job()
     job.cancel()
     # make sure the job is actually cancelled
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     assert Job.load(job.id).state == "CANCELLED"
 
 
 def test_send_signal(submit_job):
     job = submit_job()
 
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     assert Job.load(job.id).state == "RUNNING"
 
     # Send a SIGKILL (basically cancelling the Job)
     job.send_signal(9)
 
     # make sure the job is actually cancelled
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     assert Job.load(job.id).state == "CANCELLED"
 
 
 def test_suspend_unsuspend(submit_job):
     job = submit_job()
 
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     job.suspend()
     assert Job.load(job.id).state == "SUSPENDED"
 
     job.unsuspend()
     # make sure the job is actually running again
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     assert Job.load(job.id).state == "RUNNING"
 
 
@@ -121,7 +122,7 @@ def test_requeue(submit_job):
 
     assert job.requeue_count == 0
 
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     job.requeue()
     job = Job.load(job.id)
 
@@ -130,7 +131,7 @@ def test_requeue(submit_job):
 
 def test_notify(submit_job):
     job = submit_job()
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
 
     # Could check the logfile, but we just assume for now
     # that when this function raises no Exception, everything worked.
@@ -155,6 +156,35 @@ def test_get_job_queue(submit_job):
         assert isinstance(jobs[job.id], Job)
 
 
+def test_load_steps(submit_job):
+    job_list = [submit_job() for i in range(3)]
+    util.wait()
+
+    jobs = Jobs.load()
+    jobs.load_steps()
+
+    for _job in job_list:
+        job = jobs[_job.id]
+        assert job.state == "RUNNING"
+        assert job.steps
+        assert isinstance(job.steps, pyslurm.JobSteps)
+        assert job.steps.get("batch")
+
+
+def test_to_json(submit_job):
+    job_list = [submit_job() for i in range(3)]
+    util.wait()
+
+    jobs = Jobs.load()
+    jobs.load_steps()
+
+    json_data = jobs.to_json()
+    dict_data = json.loads(json_data)
+    assert dict_data
+    assert json_data
+    assert len(dict_data) >= 3
+
+
 def test_get_resource_layout_per_node(submit_job):
     # TODO
     assert True

tests/integration/test_job_steps.py

@@ -57,7 +57,7 @@ def test_load(submit_job):
 
     # Load the step info, waiting one second to make sure the Step
     # actually exists.
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     step = JobStep.load(job.id, "batch")
 
     assert step.id == "batch"
@@ -101,7 +101,7 @@ def test_load(submit_job):
 def test_collection(submit_job):
     job = submit_job(script=create_job_script_multi_step())
 
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     steps = JobSteps.load(job)
 
     assert steps
@@ -115,7 +115,7 @@ def test_collection(submit_job):
 def test_cancel(submit_job):
     job = submit_job(script=create_job_script_multi_step())
 
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     steps = JobSteps.load(job)
     assert len(steps) == 3
     assert ("batch" in steps and
@@ -124,7 +124,7 @@ def test_cancel(submit_job):
 
     steps[0].cancel()
 
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     steps = JobSteps.load(job)
     assert len(steps) == 2
     assert ("batch" in steps and
@@ -135,7 +135,7 @@ def test_modify(submit_job):
     steps = "srun -t 20 sleep 100"
     job = submit_job(script=create_job_script_multi_step(steps))
 
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     step = JobStep.load(job, 0)
     assert step.time_limit == 20
 
@@ -150,7 +150,7 @@ def test_send_signal(submit_job):
     steps = "srun -t 10 sleep 100"
     job = submit_job(script=create_job_script_multi_step(steps))
 
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     step = JobStep.load(job, 0)
     assert step.state == "RUNNING"
 
@@ -159,7 +159,7 @@ def test_send_signal(submit_job):
 
     # Make sure the job is actually cancelled.
     # If a RPCError is raised, this means the Step got cancelled.
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     with pytest.raises(RPCError):
         step = JobStep.load(job, 0)
 
@@ -173,5 +173,5 @@ def test_load_with_wrong_step_id(submit_job):
 
 def test_parse_all(submit_job):
     job = submit_job()
-    time.sleep(util.WAIT_SECS_SLURMCTLD)
+    util.wait()
     JobStep.load(job, "batch").to_dict()

tests/integration/test_node.py

@@ -20,11 +20,9 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 """test_node.py - Test the node api functions."""
 
-import sys
-import time
 import pytest
 import pyslurm
-import os
+import json
 from pyslurm import Node, Nodes, RPCError
 
 
@@ -51,10 +49,6 @@ def test_create():
         Node("testhostpyslurm2").create("idle")
 
 
-# def test_delete():
-#    node = Node("testhost1").delete()
-
-
 def test_modify():
     _, node = Nodes.load().popitem()
 
@@ -71,3 +65,13 @@ def test_modify():
 def test_parse_all():
     _, node = Nodes.load().popitem()
     assert node.to_dict()
+
+
+def test_to_json():
+    nodes = Nodes.load()
+    json_data = nodes.to_json()
+    dict_data = json.loads(json_data)
+
+    assert dict_data
+    assert len(dict_data) >= 1
+    assert json_data

tests/integration/test_partition.py

@@ -23,6 +23,7 @@
 
 import pytest
 import pyslurm
+import json
 import util
 from pyslurm import Partition, Partitions, RPCError
 
@@ -72,6 +73,16 @@ def test_parse_all():
     assert part.to_dict()
 
 
+def test_to_json():
+    parts = Partitions.load()
+    json_data = parts.to_json()
+    dict_data = json.loads(json_data)
+
+    assert dict_data
+    assert len(dict_data) >= 1
+    assert json_data
+
+
 def test_reload():
     _partnames = [util.randstr() for i in range(3)]
     _tmp_parts = Partitions(_partnames)