Merge branch '2411-check-copyright'

[arvados.git] / services / nodemanager / tests / test_jobqueue.py

diff --git a/services/nodemanager/tests/test_jobqueue.py b/services/nodemanager/tests/test_jobqueue.py
index 2c430526129e526115f4c71270a0ae90e80ea81d..8aa0835aca395af7594659e6c836e6245d189494 100644 (file)
--- a/services/nodemanager/tests/test_jobqueue.py
+++ b/services/nodemanager/tests/test_jobqueue.py
@@ -1,4 +1,7 @@
 #!/usr/bin/env python
+# Copyright (C) The Arvados Authors. All rights reserved.
+#
+# SPDX-License-Identifier: AGPL-3.0
 
 from __future__ import absolute_import, print_function
 
@@ -35,6 +38,15 @@ class ServerCalculatorTestCase(unittest.TestCase):
         servlist = self.calculate(servcalc, {'min_ram_mb_per_node': 121})
         self.assertEqual(1, len(servlist))
 
+    def test_custom_node_mem_scaling_factor(self):
+        # Simulate a custom 'node_mem_scaling' config parameter by passing
+        # the value to ServerCalculator
+        servcalc = self.make_calculator([1], node_mem_scaling=0.5)
+        servlist = self.calculate(servcalc, {'min_ram_mb_per_node': 128})
+        self.assertEqual(0, len(servlist))
+        servlist = self.calculate(servcalc, {'min_ram_mb_per_node': 64})
+        self.assertEqual(1, len(servlist))
+
     def test_implicit_server_count(self):
         servcalc = self.make_calculator([1])
         servlist = self.calculate(servcalc, {}, {'min_nodes': 3})
@@ -53,7 +65,7 @@ class ServerCalculatorTestCase(unittest.TestCase):
                                   {'min_ram_mb_per_node': 256},
                                   {'min_nodes': 6},
                                   {'min_nodes': 12},
-                                  {'min_scratch_mb_per_node': 200})
+                                  {'min_scratch_mb_per_node': 300000})
         self.assertEqual(6, len(servlist))
 
     def test_ignore_too_expensive_jobs(self):
@@ -133,7 +145,7 @@ class JobQueueMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
     def test_subscribers_get_server_lists(self, mock_squeue):
         mock_squeue.return_value = ""
 
-        self.build_monitor([{'items': [1, 2]}], self.MockCalculator())
+        self.build_monitor([{'items': [1, 2]}], self.MockCalculator(), True, True)
         self.monitor.subscribe(self.subscriber).get(self.TIMEOUT)
         self.stop_proxy(self.monitor)
         self.subscriber.assert_called_with([testutil.MockSize(1),
@@ -141,17 +153,49 @@ class JobQueueMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
 
     @mock.patch("subprocess.check_output")
     def test_squeue_server_list(self, mock_squeue):
-        mock_squeue.return_value = """1 0 0 Resources zzzzz-zzzzz-zzzzzzzzzzzzzzy
-2 0 0 Resources zzzzz-zzzzz-zzzzzzzzzzzzzzz
+        mock_squeue.return_value = """1|1024|0|Resources|zzzzz-zzzzz-zzzzzzzzzzzzzzy
+2|1024|0|Resources|zzzzz-zzzzz-zzzzzzzzzzzzzzz
+"""
+
+        super(JobQueueMonitorActorTestCase, self).build_monitor(jobqueue.ServerCalculator(
+            [(testutil.MockSize(n), {'cores': n, 'ram': n*1024, 'scratch': n}) for n in range(1, 3)]),
+                                                                True, True)
+        self.monitor.subscribe(self.subscriber).get(self.TIMEOUT)
+        self.stop_proxy(self.monitor)
+        self.subscriber.assert_called_with([testutil.MockSize(1),
+                                            testutil.MockSize(2)])
+
+    @mock.patch("subprocess.check_output")
+    def test_squeue_server_list_suffix(self, mock_squeue):
+        mock_squeue.return_value = """1|1024M|0|ReqNodeNotAvail, UnavailableNod|zzzzz-zzzzz-zzzzzzzzzzzzzzy
+1|2G|0|ReqNodeNotAvail, UnavailableNod|zzzzz-zzzzz-zzzzzzzzzzzzzzz
 """
 
         super(JobQueueMonitorActorTestCase, self).build_monitor(jobqueue.ServerCalculator(
-            [(testutil.MockSize(n), {'cores': n, 'ram': n, 'scratch': n}) for n in range(1, 3)]))
+            [(testutil.MockSize(n), {'cores': n, 'ram': n*1024, 'scratch': n}) for n in range(1, 3)]),
+                                                                True, True)
         self.monitor.subscribe(self.subscriber).get(self.TIMEOUT)
         self.stop_proxy(self.monitor)
         self.subscriber.assert_called_with([testutil.MockSize(1),
                                             testutil.MockSize(2)])
 
+    def test_coerce_to_mb(self):
+        self.assertEqual(1, jobqueue.JobQueueMonitorActor.coerce_to_mb("1"))
+        self.assertEqual(512, jobqueue.JobQueueMonitorActor.coerce_to_mb("512"))
+        self.assertEqual(512, jobqueue.JobQueueMonitorActor.coerce_to_mb("512M"))
+        self.assertEqual(1024, jobqueue.JobQueueMonitorActor.coerce_to_mb("1024M"))
+        self.assertEqual(1024, jobqueue.JobQueueMonitorActor.coerce_to_mb("1G"))
+        self.assertEqual(1536, jobqueue.JobQueueMonitorActor.coerce_to_mb("1.5G"))
+        self.assertEqual(2048, jobqueue.JobQueueMonitorActor.coerce_to_mb("2G"))
+        self.assertEqual(1025, jobqueue.JobQueueMonitorActor.coerce_to_mb("1025M"))
+        self.assertEqual(1048576, jobqueue.JobQueueMonitorActor.coerce_to_mb("1T"))
+        self.assertEqual(1572864, jobqueue.JobQueueMonitorActor.coerce_to_mb("1.5T"))
+        self.assertEqual(1073741824, jobqueue.JobQueueMonitorActor.coerce_to_mb("1P"))
+        self.assertEqual(1610612736, jobqueue.JobQueueMonitorActor.coerce_to_mb("1.5P"))
+        self.assertEqual(0, jobqueue.JobQueueMonitorActor.coerce_to_mb("0"))
+        self.assertEqual(0, jobqueue.JobQueueMonitorActor.coerce_to_mb("0M"))
+        self.assertEqual(0, jobqueue.JobQueueMonitorActor.coerce_to_mb("0G"))
+
 
 if __name__ == '__main__':
     unittest.main()