uncloud/ucloud/scheduler/tests/test_dead_host_mechanism.py

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()