# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com


import threading
import time

import hypothesis.strategies as st
import pytest
from hypothesis import given, settings

from czc.concurrency import async_map


@pytest.mark.unit
@given(st.lists(st.integers()))
@settings(max_examples=100, deadline=1000.00)
def test_async_map(inputs):
    outputs = async_map(lambda x: x, [[x] for x in inputs])
    assert outputs == inputs


@pytest.mark.unit
@given(st.lists(st.integers(), min_size=6, max_size=6))
@settings(max_examples=1, deadline=3000.00)
def test_async_map_slow_task(inputs):
    def slow_task(x):
        time.sleep(1)
        return x

    outputs = async_map(slow_task, [[x] for x in inputs])
    assert outputs == inputs


@pytest.mark.unit
@pytest.mark.parametrize(
    'tasks_run,max_concurrent',
    [
        pytest.param(6, 1, id='single_concurrent_task'),
        pytest.param(6, 2, id='dual_concurrent_tasks'),
        pytest.param(6, 10, id='large_limit_larger_than_concurrency_limit'),
        pytest.param(25, 10, id='large_limit_smaller_than_concurrency_limit'),
    ],
)
def test_max_concurrent_limits_tasks(tasks_run, max_concurrent):
    state = {'running': 0, 'tasks': [], 'snapshot': []}
    lock = threading.Lock()

    def track_concurrent_task(x):
        with lock:
            state['running'] += 1
            state['snapshot'].append(state['running'])
        time.sleep(0.5)
        with lock:
            state['running'] -= 1
            state['tasks'].append(x)

        return x

    inputs = [[x] for x in range(tasks_run)]
    async_map(track_concurrent_task, inputs, max_concurrent=max_concurrent)

    assert state['running'] == 0
    assert set(state['tasks']) == set(range(tasks_run))
    assert max(state['snapshot']) <= max_concurrent


@pytest.mark.unit
@given(st.integers(max_value=0))
def test_max_concurrency_cannot_be_less_than_1(max_concurrent):
    def task(x):
        time.sleep(0.1)
        return x

    inputs = [[x] for x in range(6)]
    with pytest.raises(ValueError):
        async_map(task, inputs, max_concurrent=max_concurrent)
