#!/usr/bin/env python3
"""CLI v2 (aws cli) в нагрузке: 50 циклов send/receive/delete с замером латентности.
Запуск на ВМ: python3 tests/cli_load_test.py"""
import boto3
import subprocess
import time
import os
import re

ENV = dict(os.environ)
EP = 'https://sqs.containerk8s.dev.nubes.ru'
sqs = boto3.client('sqs', endpoint_url=EP, region_name='us-east-1')


def cli(*a, tmo=20):
    t = time.time()
    p = subprocess.run(['aws', '--endpoint-url', EP, '--region', 'us-east-1', 'sqs', *a],
                       env=ENV, capture_output=True, text=True, timeout=tmo)
    return p, time.time() - t


def main():
    q = 'cliload-%d' % int(time.time())
    sqs.create_queue(QueueName=q)
    u = EP + '/' + q
    lat = {'send': [], 'receive': [], 'delete': []}
    failed = 0
    for i in range(50):
        p, t = cli('send-message', '--queue-url', u, '--message-body', 'msg-%d' % i)
        lat['send'].append(t)
        if p.returncode != 0:
            failed += 1
            print('send FAIL', i, p.stderr[:120])
            continue
        p, t = cli('receive-message', '--queue-url', u, '--max-number-of-messages', '1', '--visibility-timeout', '3')
        lat['receive'].append(t)
        m = re.search(r'"ReceiptHandle":\s*"([^"]+)"', p.stdout)
        if not m:
            failed += 1
            print('receive FAIL', i, p.stdout[:100], p.stderr[:100])
            continue
        p, t = cli('delete-message', '--queue-url', u, '--receipt-handle', m.group(1))
        lat['delete'].append(t)
        if p.returncode != 0:
            failed += 1
            print('delete FAIL', i, p.stderr[:120])
    for k in lat:
        v = sorted(lat[k])
        n = len(v)
        if n:
            print('%s: n=%d p50=%.0fms p95=%.0fms max=%.0fms' % (k, n, v[n // 2] * 1000, v[int(n * .95)] * 1000, v[-1] * 1000))
        else:
            print('%s: n=0' % k)
    sqs.delete_queue(QueueUrl=u)
    print('CLI load done, failed=%d' % failed)


if __name__ == '__main__':
    main()