Fix RepresenterError in db slurm new assoc (#37)

* bugfix: the new asosc command was raising a ruamel serialization error while attempting to display the newly created qos

* v1.3.3

Author: camillescottatwork

cheeto/__init__.py

@@ -1 +1 @@
-__version__ = '1.3.2'
+__version__ = '1.3.3'

cheeto/cmds/database.py

@@ -1190,7 +1190,18 @@ def cmd_slurm_new_assoc(args: Namespace):
         console.print(f'[red] QOS {args.qos} does not exist.')
         return ExitCode.DOES_NOT_EXIST
 
-    console.print(assoc.pretty())
+    output = dumps_yaml(_show_slurm_assoc(assoc))
+    console.print(highlight_yaml(output))
+
+
+def _show_slurm_assoc(assoc: SiteSlurmAssociation) -> dict:
+    return assoc._pretty(
+        lift=['partition'],
+        skip=('id',),
+        formatters={
+            'group': '{data.groupname}',
+        }
+    )
 
 
 #########################################

cheeto/database/__init__.py

@@ -72,6 +72,7 @@
     query_user_access,
     query_user_status,
     query_group_exists,
+    query_slurm_association,
     tag_comment,
     add_user_comment,
     set_user_shell,

cheeto/database/crud.py

@@ -882,6 +882,13 @@ def create_slurm_association(sitename: str, partitionname: str, groupname: str,
     return assoc
 
 
+def query_slurm_association(sitename: str, qosname: str, partition: str, groupname: str):
+    qos = SiteSlurmQOS.objects.get(qosname=qosname, sitename=sitename)
+    group = SiteGroup.objects.get(groupname=groupname, sitename=sitename)
+    partition = SiteSlurmPartition.objects.get(partitionname=partition, sitename=sitename)
+    return SiteSlurmAssociation.objects.get(sitename=sitename, qos=qos, group=group, partition=partition)
+
+
 def slurm_qos_state(sitename: str) -> dict[str, SiteSlurmQOS]:
     return {s.qosname: s.to_puppet() for s in SiteSlurmQOS.objects(sitename=sitename)}
 

cheeto/tests/conftest.py

@@ -78,6 +78,19 @@ def run_shell_cmd(cmd, print_stderr=True, in_directory=None):
         os.chdir(cwd)
 
 
+@pytest.fixture
+def run_cmd(config_file):
+    def _run_cmd(*args):
+        from logging import getLogger
+        from ..cmds.__main__ import commands
+        logger = getLogger(__name__)
+        args = [str(arg) for arg in args]
+        args.extend(['--config', str(config_file)])
+        logger.info(f'running: {" ".join(args)}')
+        return commands.run(args)
+    return _run_cmd
+
+
 @pytest.fixture(scope='session', autouse=True)
 def setup_logging():
     _setup_logging(sys.stdout)

cheeto/tests/test_db.py

@@ -109,15 +109,35 @@ def test_create_qos_cleaning(self):
         assert qos.user_limits.mem == '1024M'
         assert qos.job_limits.mem == '1024M'
 
-    def test_create_qos_command(self, config_file):
-        run_shell_cmd(['cheeto', 'database', 'slurm', 'new', 'qos',
-                       '--config', config_file,
-                       '--qosname', 'test-qos',
-                       '--site', 'test-site',
-                       '--group-limits', 'cpus=16,mem=1G,gpus=0',
-                       '--user-limits', 'cpus=16,mem=1G,gpus=0',
-                       '--job-limits', 'cpus=16,mem=1G,gpus=0'])
+    def test_create_qos_command(self, run_cmd):
+        run_cmd('database', 'slurm', 'new', 'qos',
+                '--qosname', 'test-qos',
+                '--site', 'test-site',
+                '--group-limits', 'cpus=16,mem=1G,gpus=0',
+                '--user-limits', 'cpus=16,mem=1G,gpus=0',
+                '--job-limits', 'cpus=16,mem=1G,gpus=0')
         assert SiteSlurmQOS.objects.count() == 1
         assert SiteSlurmQOS.objects.get(qosname='test-qos').group_limits.mem == '1024M'
         assert SiteSlurmQOS.objects.get(qosname='test-qos').user_limits.mem == '1024M'
         assert SiteSlurmQOS.objects.get(qosname='test-qos').job_limits.mem == '1024M'
+
+    def test_create_assoc_command(self, run_cmd):
+        create_group('test-group', 10000, sites=['test-site'])
+        create_slurm_partition('test-partition', 'test-site')
+
+        run_cmd('database', 'slurm', 'new', 'qos',
+                '--qosname', 'test-qos',
+                '--site', 'test-site',
+                '--group-limits', 'cpus=16,mem=1G,gpus=0',
+                '--user-limits', 'cpus=16,mem=1G,gpus=0',
+                '--job-limits', 'cpus=16,mem=1G,gpus=0')
+        run_cmd('database', 'slurm', 'new', 'assoc',
+                '--site', 'test-site',
+                '--group', 'test-group',
+                '--partition', 'test-partition',
+                '--qos', 'test-qos')
+        assert SiteSlurmAssociation.objects.count() == 1
+        assoc = query_slurm_association('test-site', 'test-qos', 'test-partition', 'test-group')
+        assert assoc.group.groupname == 'test-group'
+        assert assoc.partition.partitionname == 'test-partition'
+        assert assoc.qos.qosname == 'test-qos'

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cheeto"
-version = "1.3.2"
+version = "1.3.3"
 description = "Utilities for the UC Davis HPC Core Facility."
 license = "Proprietary"
 authors = [