bugfix: size_to_megs was not imported in SlurmQOS clean method

Author: camillescottatwork

cheeto/database/slurm.py

@@ -19,6 +19,7 @@
 
 from ..puppet import (SlurmQOSTRES as PuppetSlurmQOSTRES,
                       SlurmQOS as PuppetSlurmQOS)
+from ..utils import size_to_megs
 
 from .base import BaseDocument
 from .fields import (POSIXNameField,

cheeto/tests/conftest.py

@@ -133,9 +133,16 @@ def start_mongodb(tmp_path_factory):
 
 
 @pytest.fixture(scope='session')
-def config(request):
+def config_file(tmp_path_factory):
     data_dir = Path(__file__).parent / 'data'
-    return get_config(data_dir / 'config.yaml')
+    config_file = tmp_path_factory.mktemp('cheeto-config') / 'config.yaml'
+    shutil.copy(data_dir / 'config.yaml', config_file)
+    return config_file
+
+
+@pytest.fixture(scope='session')
+def config(config_file):
+    return get_config(config_file)
 
 
 @pytest.fixture(scope='session')

cheeto/tests/test_db.py

@@ -5,7 +5,7 @@
 
 from ..database import *
 
-from .conftest import drop_database
+from .conftest import drop_database, run_shell_cmd
 
 
 def test_connect_to_database(db_config):
@@ -77,4 +77,47 @@ def test_add_duplicate_site_user(self):
         user, group = create_user(*self.USER_ARGS)
         add_site_user('test-site', user)
         with pytest.raises(DuplicateSiteUser):
-            add_site_user('test-site', user)
+            add_site_user('test-site', user)
+
+
+class TestSlurm:
+
+    @pytest.fixture(autouse=True)
+    def setup_site(self, db_config):
+        drop_database(db_config)
+        create_site('test-site', 'test.site.com')
+        yield
+        drop_database(db_config)
+
+    def test_create_qos(self):
+        qos = create_slurm_qos('test-qos',
+                               'test-site',
+                               group_limits=SlurmTRES(cpus=16, mem='1000M', gpus=0),
+                               user_limits=SlurmTRES(cpus=16, mem='1000M', gpus=0),
+                               job_limits=SlurmTRES(cpus=16, mem='1000M', gpus=0))
+        assert qos.qosname == 'test-qos'
+        assert qos.sitename == 'test-site'
+        assert SiteSlurmQOS.objects.count() == 1
+
+    def test_create_qos_cleaning(self):
+        qos = create_slurm_qos('test-qos',
+                               'test-site',
+                               group_limits=SlurmTRES(cpus=16, mem='1G', gpus=0),
+                               user_limits=SlurmTRES(cpus=16, mem='1G', gpus=0),
+                               job_limits=SlurmTRES(cpus=16, mem='1G', gpus=0))
+        assert qos.group_limits.mem == '1024M'
+        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'])
+        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'

cheeto/tests/test_slurm.py

@@ -0,0 +1,15 @@
+import pytest
+from ..types import parse_qos_tres
+
+
+def test_parse_qos_tres():
+    assert parse_qos_tres('mem=1000M,cpus=16,gpus=0') == {'mem': '1000M', 'cpus': '16', 'gpus': '0'}
+
+
+def test_parse_qos_tres_none():
+    assert parse_qos_tres(None) == {'mem': None, 'cpus': None, 'gpus': None}
+
+
+def test_parse_qos_tres_empty():
+    with pytest.raises(ValueError):
+        print(parse_qos_tres(''))