import concurrent.futures
import requests
import time
import statistics


def call_api(url, method, params, json, print_res):
    start = time.time()
    response = requests.request(method, url, params=params, json=json)
    if print_res:
        print(response.json())
    end = time.time()
    return (end - start) * 1000, response.status_code, len(response.content)


def test_api_tps(url, method, params=None, json=None, concurrency=1, duration=1, print_res=False):
    """
    测试 API 的 TPS 等并发数据
    :param url: url示例：baidu.com
    :param method: 请求方式： get/post
    :param params: url参数：{"id":1}
    :param json:  jsonBody参数：{"id":1}
    :param concurrency: 并发数
    :param duration: 持续时间 秒
    :param print_res: 是否打印api的返回值
    :return:
    """
    successes = 0
    failures = 0
    times = []
    data_lengths = []

    with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor:
        start = time.time()
        futures = [executor.submit(call_api, url, method, params, json, print_res) for i in range(concurrency)]

        while time.time() - start < duration:
            completed, futures = concurrent.futures.wait(futures, timeout=0)
            for future in completed:
                elapsed, status, length = future.result()
                times.append(elapsed)
                data_lengths.append(length)
                if status == 200:
                    successes += 1
                else:
                    failures += 1
                futures.add(executor.submit(call_api, url, method, params, json, print_res))
        # Wait for the remaining futures to complete
        if futures:
            completed, _ = concurrent.futures.wait(futures)
            for future in completed:
                elapsed, status, length = future.result()
                times.append(elapsed)
                data_lengths.append(length)
                if status == 200:
                    successes += 1
                else:
                    failures += 1

    total = successes + failures
    tps = total / duration
    mean = statistics.mean(times)
    min_time = min(times)
    max_time = max(times)

    print(f'TPS:{tps}; ', end=' ')
    print(f'平均时间(ms): {mean:.0f}; ', end=' ')
    print(f'最小时间(ms): {min_time:.0f}; ', end=' ')
    print(f'最大时间(ms): {max_time:.0f}; ', end=' ')
    print(f'成功数: {successes}; ', end=' ')
    print(f'失败数: {failures}; ', end=' ')
    print(f'总数据长度: {sum(data_lengths)}; ', end=' ')


if __name__ == '__main__':
    test_api_tps('https://httpbin.tianyunperfect.cn/get', 'get', params={"a": 1}, duration=2)