Efforts to make ucloud a python package

ucloud/scheduler/tests/test_dead_host_mechanism.py

@@ -0,0 +1,75 @@
+import sys
+import unittest
+from datetime import datetime
+from os.path import dirname
+
+BASE_DIR = dirname(dirname(__file__))
+sys.path.insert(0, BASE_DIR)
+
+from main import (
+    dead_host_detection,
+    dead_host_mitigation,
+    config
+)
+
+
+class TestDeadHostMechanism(unittest.TestCase):
+    def setUp(self):
+        self.client = config.etcd_client
+        self.host_prefix = "/test/host"
+        self.vm_prefix = "/test/vm"
+
+        self.client.client.delete_prefix(self.host_prefix)
+        self.client.client.delete_prefix(self.vm_prefix)
+
+        self.create_hosts()
+
+    def create_hosts(self):
+        host1 = {
+            "cpu": 32,
+            "ram": 128,
+            "hdd": 1024,
+            "sdd": 0,
+            "status": "ALIVE",
+            "last_heartbeat": datetime.utcnow().isoformat(),
+        }
+        host2 = {
+            "cpu": 16,
+            "ram": 64,
+            "hdd": 512,
+            "sdd": 0,
+            "status": "ALIVE",
+            "last_heartbeat": datetime(2011, 1, 1).isoformat(),
+        }
+
+        host3 = {"cpu": 16, "ram": 32, "hdd": 256, "sdd": 256}
+        host4 = {
+            "cpu": 16,
+            "ram": 32,
+            "hdd": 256,
+            "sdd": 256,
+            "status": "DEAD",
+            "last_heartbeat": datetime(2011, 1, 1).isoformat(),
+        }
+        with self.client.client.lock("lock"):
+            self.client.put(f"{self.host_prefix}/1", host1, value_in_json=True)
+            self.client.put(f"{self.host_prefix}/2", host2, value_in_json=True)
+            self.client.put(f"{self.host_prefix}/3", host3, value_in_json=True)
+            self.client.put(f"{self.host_prefix}/4", host4, value_in_json=True)
+
+    def test_dead_host_detection(self):
+        hosts = self.client.get_prefix(self.host_prefix, value_in_json=True)
+        deads = dead_host_detection(hosts)
+        self.assertEqual(deads, ["/test/host/2", "/test/host/3"])
+        return deads
+
+    def test_dead_host_mitigation(self):
+        deads = self.test_dead_host_detection()
+        dead_host_mitigation(self.client, deads)
+        hosts = self.client.get_prefix(self.host_prefix, value_in_json=True)
+        deads = dead_host_detection(hosts)
+        self.assertEqual(deads, [])
+
+
+if __name__ == "__main__":
+    unittest.main()