can return all jobs within a drain window that are idle. Not just no draining

src/lib/Components/system/CraySystem.py

@@ -679,7 +679,7 @@ class CraySystem(BaseSystem):
             retlist.extend(expand_num_list(locs))
         return retlist
 
-    def _assemble_queue_data(self, job, idle_only=True, no_draining=False):
+    def _assemble_queue_data(self, job, idle_only=True, drain_time=None):
         '''put together data for a queue, or queue-like reservation structure.
 
         Input:
@@ -746,9 +746,10 @@ class CraySystem(BaseSystem):
                 unavailable_nodes = [node_id for node_id in node_id_list
                         if self.nodes[str(node_id)].status in
                         self.nodes[str(node_id)].DOWN_STATUSES]
-            if no_draining:
+            if drain_time is not None:
                 unavailable_nodes.extend([node_id for node_id in node_id_list
-                    if self.nodes[str(node_id)].draining])
+                    if (self.nodes[str(node_id)].draining and
+                        self.nodes[str(node_id)].drain_until < int(drain_time))])
         for node_id in set(unavailable_nodes):
             node_id_list.remove(node_id)
         return node_id_list

testsuite/TestCobalt/TestComponents/test_cray.py

@@ -295,9 +295,30 @@ class TestCraySystem(object):
         self.system.nodes['3'].status = 'allocated'
         self.system.nodes['4'].set_drain(100, 1)
         nodelist = self.system._assemble_queue_data(self.base_job,
-                no_draining=True)
+                drain_time=150)
         assert_match(sorted(nodelist), ['5'], "Bad Nodelist")
 
+    def test_assemble_queue_data_attrs_within_draining(self):
+        '''CraySystem._assemble_queue_data: return idle and draining if within
+        time'''
+        self.system.nodes['1'].status = 'busy'
+        self.system.nodes['2'].status = 'down'
+        self.system.nodes['3'].set_drain(50.0, 2)
+        self.system.nodes['4'].set_drain(100.0, 1)
+        nodelist = self.system._assemble_queue_data(self.base_job,
+                drain_time=90.0)
+        assert_match(sorted(nodelist), ['4', '5'], "Bad Nodelist")
+
+    def test_assemble_queue_data_attrs_match_draining(self):
+        '''CraySystem._assemble_queue_data: return idle and matched drain node'''
+        self.system.nodes['1'].status = 'busy'
+        self.system.nodes['2'].status = 'down'
+        self.system.nodes['3'].status = 'allocated'
+        self.system.nodes['4'].set_drain(100.0, 1)
+        nodelist = self.system._assemble_queue_data(self.base_job,
+                drain_time=100.0)
+        assert_match(sorted(nodelist), ['4', '5'], "Bad Nodelist")
+
     def test_find_queue_equivalence_classes_single(self):
         '''CraySystem.find_queue_equivalence_classes: single queue'''
         self.system.find_queue_equivalence_classes([], ['default'], [])