#!/usr/bin/env python
import sys
import time
import requests


API_ENDPOINT = "http://127.0.0.1:8261/api/v1/cluster"


def list_master_success(master_count):
    url = API_ENDPOINT + "/masters"
    resp = requests.get(url=url)
    assert resp.status_code == 200
    data = resp.json()
    print("list_master_success resp=", data)
    assert data["total"] == int(master_count)


def delete_master_with_retry_success(master_name):
    """
    offline dm-master node is done by offline etcd member but operate etcd cluster may met `etcdserver: unhealthy cluster`
    so add some retry here
    """
    url = API_ENDPOINT + "/masters/" + master_name
    for i in range(0, 30):
        resp = requests.delete(url=url)
        if resp.status_code != 204:
            print("delete_master_failed resp=", resp.json(), "retry cnt=", i)
            time.sleep(1)
        else:
            assert resp.status_code == 204
            print("delete_master_with_retry_success")
            return
    raise Exception("delete_master_with_retry_success failed")


def list_worker_success(worker_count):
    url = API_ENDPOINT + "/workers"
    resp = requests.get(url=url)
    assert resp.status_code == 200
    data = resp.json()
    print("list_worker_success resp=", data)
    assert data["total"] == int(worker_count)


def delete_worker_failed(worker_name):
    url = API_ENDPOINT + "/workers/" + worker_name
    resp = requests.delete(url=url)
    assert resp.status_code == 400
    print("delete_worker_failed resp=", resp.json())


def delete_worker_with_retry_success(worker_name):
    """
    offline dm-worker node is done by offline etcd member but operate etcd cluster may met `etcdserver: unhealthy cluster`
    so add some retry here
    """
    url = API_ENDPOINT + "/workers/" + worker_name
    for i in range(0, 30):
        resp = requests.delete(url=url)
        if resp.status_code != 204:
            print("delete_worker_failed resp=", resp.json(), "retry cnt=", i)
            time.sleep(1)
        else:
            assert resp.status_code == 204
            print("delete_worker_with_retry_success")
            return
    raise Exception("delete_master_with_retry_success failed")


if __name__ == "__main__":
    FUNC_MAP = {
        "list_master_success": list_master_success,
        "delete_master_with_retry_success": delete_master_with_retry_success,
        "list_worker_success": list_worker_success,
        "delete_worker_failed": delete_worker_failed,
        "delete_worker_with_retry_success": delete_worker_with_retry_success,
    }

    func = FUNC_MAP[sys.argv[1]]
    if len(sys.argv) >= 2:
        func(*sys.argv[2:])
    else:
        func()