Handle allocation while machine type not specify

confluentinc · Dec 10, 2019 · ab512cc · ab512cc
1 parent 0076773
commit ab512cc
Showing 1 changed file with 9 additions and 14 deletions.
diff --git a/ducktape/cluster/node_container.py b/ducktape/cluster/node_container.py
@@ -14,6 +14,7 @@
 
 from six import iteritems
 from operator import attrgetter
+from ducktape.cluster.remoteaccount import MachineType
 
 
 class NodeNotPresentError(Exception):
@@ -183,11 +184,10 @@ def satisfy(avail_node, req_node):
         Return true if available node satisfies the minimum requirement of requested node.
         """
         if req_node.machine_type is None:
-            return True
-        elif avail_node.machine_type is None or \
-                avail_node.machine_type.cpu_core < req_node.machine_type.cpu_core or \
-                avail_node.machine_type.mem_size_gb < req_node.machine_type.mem_size_gb or \
-                avail_node.machine_type.disk_size_gb < req_node.machine_type.disk_size_gb:
+            req_node.machine_type = MachineType()
+        if avail_node.machine_type.cpu_core < req_node.machine_type.cpu_core or \
+           avail_node.machine_type.mem_size_gb < req_node.machine_type.mem_size_gb or \
+           avail_node.machine_type.disk_size_gb < req_node.machine_type.disk_size_gb:
             return False
         for d_name, d_size in req_node.machine_type.additional_disks.items():
             if avail_node.machine_type.additional_disks.get(d_name, 0) < d_size:
@@ -197,19 +197,14 @@ def satisfy(avail_node, req_node):
     @staticmethod
     def sort(nodes, reverse=False):
         """
-        Return sorted list of nodes based on machine_type. If any the machine
+        Return sorted list of nodes based on machine_type.
         """
-        sorted_nodes = []
-        type_based_nodes = []
         for node in nodes:
             if node.machine_type is None:
-                sorted_nodes.append(node)
-            else:
-                type_based_nodes.append(node)
+                node.machine_type = MachineType()
 
-        sorted_nodes.extend(sorted(type_based_nodes, key=attrgetter('machine_type.cpu_core', 'machine_type.mem_size_gb',
-                                                                    'machine_type.disk_size_gb',
-                                                                    'machine_type.additional_disks')))
+        sorted_nodes = sorted(nodes, key=attrgetter('machine_type.cpu_core', 'machine_type.mem_size_gb',
+                                                    'machine_type.disk_size_gb', 'machine_type.additional_disks'))
         return list(reversed(sorted_nodes)) if reverse else sorted_nodes
 
     def clone(self):