Fix unit tests

tests/cluster/check_cluster.py

@@ -16,20 +16,21 @@
 
 from ducktape.cluster.cluster import ClusterNode
 from ducktape.cluster.cluster_spec import NodeSpec, ClusterSpec, LINUX, WINDOWS
+from ducktape.cluster.remoteaccount import MachineType
 from tests.ducktape_mock import FakeCluster
 
-FakeRemoteAccount = collections.namedtuple('FakeRemoteAccount', ['operating_system'])
+FakeRemoteAccount = collections.namedtuple('FakeRemoteAccount', ['operating_system', 'machine_type'])
 
 
 class CheckCluster(object):
 
     def setup_method(self, _):
         self.cluster = FakeCluster(0)
-        self.cluster._available_nodes.add_node(ClusterNode(FakeRemoteAccount(operating_system=LINUX)))
-        self.cluster._available_nodes.add_node(ClusterNode(FakeRemoteAccount(operating_system=LINUX)))
-        self.cluster._available_nodes.add_node(ClusterNode(FakeRemoteAccount(operating_system=WINDOWS)))
-        self.cluster._available_nodes.add_node(ClusterNode(FakeRemoteAccount(operating_system=WINDOWS)))
-        self.cluster._available_nodes.add_node(ClusterNode(FakeRemoteAccount(operating_system=WINDOWS)))
+        self.cluster._available_nodes.add_node(ClusterNode(FakeRemoteAccount(operating_system=LINUX, machine_type=MachineType())))
+        self.cluster._available_nodes.add_node(ClusterNode(FakeRemoteAccount(operating_system=LINUX, machine_type=MachineType())))
+        self.cluster._available_nodes.add_node(ClusterNode(FakeRemoteAccount(operating_system=WINDOWS, machine_type=MachineType())))
+        self.cluster._available_nodes.add_node(ClusterNode(FakeRemoteAccount(operating_system=WINDOWS, machine_type=MachineType())))
+        self.cluster._available_nodes.add_node(ClusterNode(FakeRemoteAccount(operating_system=WINDOWS, machine_type=MachineType())))
 
     def spec(self, linux_nodes, windows_nodes):
         nodes = []

tests/cluster/check_cluster_spec.py

@@ -25,4 +25,23 @@ def check_to_string(self):
         empty = ClusterSpec.empty()
         assert "[]" == str(empty)
         simple_linux_5 = ClusterSpec.simple_linux(5)
-        assert '[{"num_nodes": 5, "os": "linux"}]' == str(simple_linux_5)
+        assert '[{"additional_disks(GB)": {}, "cpu": 0, "disk(GB)": 0, "mem(GB)": 0, "num_nodes": 5, "os": "linux"}]' \
+               == str(simple_linux_5)
+
+    def check_from_dict(self):
+        empty = ClusterSpec.empty()
+        assert "[]" == str(empty)
+        node_specs_dict = {'cpu': 2, 'mem': '2GB', 'disk': '30GB', 'num_nodes': 2}
+        custom_linux_1 = ClusterSpec.from_dict(node_specs_dict)
+        assert '[{"additional_disks(GB)": {}, "cpu": 2, "disk(GB)": 30.0, "mem(GB)": 2.0, "num_nodes": 2, "os": "linux"}]' \
+               == str(custom_linux_1)
+
+    def check_from_list(self):
+        empty = ClusterSpec.empty()
+        assert "[]" == str(empty)
+        node_specs_dict_list = [{'cpu': 2, 'mem': '2GB', 'disk': '20GB', 'num_nodes': 2},
+                                {'cpu': 4, 'mem': '4GB', 'disk': '40GB', 'num_nodes': 4}]
+        custom_linux_2 = ClusterSpec.from_list(node_specs_dict_list)
+        assert '[{"additional_disks(GB)": {}, "cpu": 2, "disk(GB)": 20.0, "mem(GB)": 2.0, "num_nodes": 2, "os": "linux"},' \
+               ' {"additional_disks(GB)": {}, "cpu": 4, "disk(GB)": 40.0, "mem(GB)": 4.0, "num_nodes": 4, "os": "linux"}]' \
+               == str(custom_linux_2)

tests/cluster/check_finite_subcluster.py

@@ -15,6 +15,7 @@
 from ducktape.cluster.cluster_spec import LINUX
 from ducktape.cluster.finite_subcluster import FiniteSubcluster
 from ducktape.cluster.node_container import InsufficientResourcesError, NodeNotPresentError
+from ducktape.cluster.remoteaccount import MachineType
 from ducktape.services.service import Service
 import pickle
 import pytest
@@ -25,6 +26,9 @@ class MockFiniteSubclusterNode:
     def operating_system(self):
         return LINUX
 
+    @property
+    def machine_type(self):
+        return MachineType()
 
 class CheckFiniteSubcluster(object):
     single_node_cluster_json = {"nodes": [{"hostname": "localhost"}]}

tests/cluster/check_json.py

@@ -31,42 +31,48 @@ class CheckJsonCluster(object):
     def check_invalid_json(self):
         # Missing list of nodes
         with pytest.raises(ValueError):
-            JsonCluster({})
+            JsonCluster(cluster_json={}, is_type_based=False)
 
         # Missing hostname, which is required
         with pytest.raises(ValueError):
-            JsonCluster({"nodes": [{}]})
+            JsonCluster(cluster_json={"nodes": [{}]}, is_type_based=False)
 
     @staticmethod
     def cluster_hostnames(nodes):
         return set([node.account.hostname for node in nodes])
 
     def check_cluster_size(self):
-        cluster = JsonCluster({"nodes": []})
+        cluster = JsonCluster(cluster_json={"nodes": []}, is_type_based=False)
         assert len(cluster) == 0
 
         n = 10
         cluster = JsonCluster(
-            {"nodes": [
-                {"ssh_config": {"hostname": "localhost%d" % x}} for x in range(n)]})
+            cluster_json={
+                "nodes": [
+                    {"ssh_config": {"host": "localhost%d" % x}} for x in range(n)]},
+            is_type_based=False)
 
         assert len(cluster) == n
 
     def check_pickleable(self):
         cluster = JsonCluster(
-            {"nodes": [
-                {"ssh_config": {"host": "localhost1"}},
-                {"ssh_config": {"host": "localhost2"}},
-                {"ssh_config": {"host": "localhost3"}}]})
+            cluster_json={
+                "nodes": [
+                    {"ssh_config": {"host": "localhost1"}},
+                    {"ssh_config": {"host": "localhost2"}},
+                    {"ssh_config": {"host": "localhost3"}}]},
+            is_type_based=False)
 
         pickle.dumps(cluster)
 
     def check_allocate_free(self):
         cluster = JsonCluster(
-            {"nodes": [
-                {"ssh_config": {"host": "localhost1"}},
-                {"ssh_config": {"host": "localhost2"}},
-                {"ssh_config": {"host": "localhost3"}}]})
+            cluster_json = {
+                "nodes": [
+                    {"ssh_config": {"host": "localhost1"}},
+                    {"ssh_config": {"host": "localhost2"}},
+                    {"ssh_config": {"host": "localhost3"}}]},
+            is_type_based=False)
 
         assert len(cluster) == 3
         assert(cluster.num_available_nodes() == 3)
@@ -93,9 +99,10 @@ def check_parsing(self):
         """ Checks that RemoteAccounts are generated correctly from input JSON"""
 
         node = JsonCluster(
-            {
+            cluster_json={
                 "nodes": [
-                    {"ssh_config": {"host": "hostname"}}]}).alloc(Service.setup_cluster_spec(num_nodes=1))[0]
+                    {"ssh_config": {"host": "hostname"}}]},
+            is_type_based=False).alloc(Service.setup_cluster_spec(num_nodes=1))[0]
 
         assert node.account.hostname == "hostname"
         assert node.account.user is None
@@ -106,9 +113,12 @@ def check_parsing(self):
             "hostname": "localhost",
             "port": 22
         }
-        node = JsonCluster({"nodes": [{"hostname": "hostname",
+        node = JsonCluster(
+            cluster_json={
+                "nodes": [{"hostname": "hostname",
                                        "user": "user",
-                                       "ssh_config": ssh_config}]}).alloc(Service.setup_cluster_spec(num_nodes=1))[0]
+                                       "ssh_config": ssh_config}]}, 
+            is_type_based=False).alloc(Service.setup_cluster_spec(num_nodes=1))[0]
 
         assert node.account.hostname == "hostname"
         assert node.account.user == "user"
@@ -120,16 +130,18 @@ def check_parsing(self):
         assert node.account.ssh_config.port == 22
 
     def check_exhausts_supply(self):
-        cluster = JsonCluster(self.single_node_cluster_json)
+        cluster = JsonCluster(cluster_json=self.single_node_cluster_json, is_type_based=False)
         with pytest.raises(InsufficientResourcesError):
             cluster.alloc(Service.setup_cluster_spec(num_nodes=2))
 
     def check_node_names(self):
         cluster = JsonCluster(
-            {"nodes": [
-                {"ssh_config": {"host": "localhost1"}},
-                {"ssh_config": {"host": "localhost2"}},
-                {"ssh_config": {"host": "localhost3"}}]})
+            cluster_json={
+                "nodes": [
+                    {"ssh_config": {"host": "localhost1"}},
+                    {"ssh_config": {"host": "localhost2"}},
+                    {"ssh_config": {"host": "localhost3"}}]},
+            is_type_based=False)
         hosts = set(["localhost1", "localhost2", "localhost3"])
         nodes = cluster.alloc(cluster.available())
         assert hosts == set(node.name for node in nodes)

tests/cluster/check_localhost.py

@@ -20,13 +20,13 @@
 
 class CheckLocalhostCluster(object):
     def setup_method(self, _):
-        self.cluster = LocalhostCluster()
+        self.cluster = LocalhostCluster(is_type_based=False)
 
     def check_size(self):
         len(self.cluster) >= 2 ** 31 - 1
 
     def check_pickleable(self):
-        cluster = LocalhostCluster()
+        cluster = LocalhostCluster(is_type_based=False)
         pickle.dumps(cluster)
 
     def check_request_free(self):

tests/cluster/check_remoteaccount.py

@@ -100,7 +100,8 @@ def check_remote_account_equality(self):
         kwargs = {
             "ssh_config": ssh_config,
             "externally_routable_ip": "345",
-            "logger": logging.getLogger(__name__)
+            "logger": logging.getLogger(__name__),
+            "is_type_based": False
         }
         r1 = RemoteAccount(**kwargs)
         r2 = RemoteAccount(**kwargs)

tests/cluster/check_vagrant.py

@@ -66,7 +66,7 @@ def _set_monkeypatch_attr(self, monkeypatch):
 
     def check_pickleable(self, monkeypatch):
         self._set_monkeypatch_attr(monkeypatch)
-        cluster = VagrantCluster()
+        cluster = VagrantCluster(is_type_based=False)
         pickle.dumps(cluster)
 
     def check_one_host_parsing(self, monkeypatch):
@@ -75,7 +75,7 @@ def check_one_host_parsing(self, monkeypatch):
         """
         self._set_monkeypatch_attr(monkeypatch)
 
-        cluster = VagrantCluster()
+        cluster = VagrantCluster(is_type_based=False)
         assert len(cluster) == 2
         assert cluster.num_available_nodes() == 2
         node1, node2 = cluster.alloc(Service.setup_cluster_spec(num_nodes=2))
@@ -96,7 +96,7 @@ def check_cluster_file_write(self, monkeypatch):
         self._set_monkeypatch_attr(monkeypatch)
         assert not os.path.exists(self.cluster_file)
 
-        cluster = VagrantCluster(cluster_file=self.cluster_file)
+        cluster = VagrantCluster(cluster_file=self.cluster_file, is_type_based=False)
         cluster_json_expected = {}
         nodes = [
             {
@@ -159,7 +159,7 @@ def check_cluster_file_read(self, monkeypatch):
                   indent=2, separators=(',', ': '), sort_keys=True)
 
         # Load the cluster from the json file we just created
-        cluster = VagrantCluster(cluster_file=self.cluster_file)
+        cluster = VagrantCluster(cluster_file=self.cluster_file, is_type_based=False)
 
         assert len(cluster) == 2
         assert cluster.num_available_nodes() == 2

tests/ducktape_mock.py

@@ -18,7 +18,7 @@
 from ducktape.tests.session import SessionContext
 from ducktape.tests.test import TestContext
 from ducktape.cluster.linux_remoteaccount import LinuxRemoteAccount
-from ducktape.cluster.remoteaccount import RemoteAccountSSHConfig
+from ducktape.cluster.remoteaccount import RemoteAccountSSHConfig, MachineType
 from mock import MagicMock
 
 
@@ -35,6 +35,9 @@ class FakeClusterNode(object):
     def operating_system(self):
         return LINUX
 
+    @property
+    def machine_type(self):
+        return MachineType()
 
 class FakeCluster(Cluster):
     """A cluster class with counters, but no actual node objects"""
@@ -95,4 +98,4 @@ def __init__(self):
             hostname="localhost",
             port=22)
 
-        super(MockAccount, self).__init__(ssh_config, externally_routable_ip="localhost", logger=None)
+        super(MockAccount, self).__init__(ssh_config, externally_routable_ip="localhost", is_type_based=False, logger=None)

tests/runner/check_runner.py

@@ -65,7 +65,7 @@ def _do_expand(self, test_file, test_class, test_methods, cluster=None, session_
 
     def check_simple_run(self):
         """Check expected behavior when running a single test."""
-        mock_cluster = LocalhostCluster(num_nodes=1000)
+        mock_cluster = LocalhostCluster(num_nodes=1000, is_type_based=False)
         session_context = tests.ducktape_mock.session_context()
 
         test_methods = [TestThingy.test_pi, TestThingy.test_ignore1, TestThingy.test_ignore2]
@@ -86,7 +86,7 @@ def check_exit_first(self):
         """Confirm that exit_first in session context has desired effect of preventing any tests from running
         after the first test failure.
         """
-        mock_cluster = LocalhostCluster(num_nodes=1000)
+        mock_cluster = LocalhostCluster(num_nodes=1000, is_type_based=False)
         session_context = tests.ducktape_mock.session_context(**{"exit_first": True})
 
         test_methods = [FailingTest.test_fail]
@@ -100,7 +100,7 @@ def check_exit_first(self):
     def check_exits_if_failed_to_initialize(self):
         """Validate that runner exits correctly when tests failed to initialize.
         """
-        mock_cluster = LocalhostCluster(num_nodes=1000)
+        mock_cluster = LocalhostCluster(num_nodes=1000, is_type_based=False)
         session_context = tests.ducktape_mock.session_context()
 
         ctx_list = self._do_expand(test_file=FAILS_TO_INIT_TEST_FILE, test_class=FailsToInitTest,

tests/runner/check_runner_memory.py

@@ -62,7 +62,7 @@ def _run_single_test(self, test_context):
 
 class CheckMemoryUsage(object):
     def setup_method(self, _):
-        self.cluster = LocalhostCluster(num_nodes=100)
+        self.cluster = LocalhostCluster(num_nodes=100, is_type_based=False)
         self.session_context = tests.ducktape_mock.session_context()
 
     def check_for_inter_test_memory_leak(self):

tests/services/check_service.py

@@ -40,7 +40,7 @@ def idx(self, node):
 class CheckAllocateFree(object):
 
     def setup_method(self, _):
-        self.cluster = LocalhostCluster()
+        self.cluster = LocalhostCluster(is_type_based=False)
         self.session_context = session_context()
         self.context = test_context(self.session_context, cluster=self.cluster)
 

tests/tests/check_session.py

@@ -40,7 +40,7 @@ def check_pickleable(self):
         kwargs = {
             "session_id": "hello-123",
             "results_dir": self.tempdir,
-            "cluster": LocalhostCluster(),
+            "cluster": LocalhostCluster(is_type_based=False),
             "globals": {}
         }
         session_context = SessionContext(**kwargs)