# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import time
from collections import deque

import pytest

from nvflare.apis.event_type import EventType
from nvflare.apis.fl_context import FLContext, FLContextManager
from nvflare.app_common.resource_managers.list_resource_manager import ListResourceManager


class MockEngine:
    def __init__(self, run_name="exp1"):
        self.fl_ctx_mgr = FLContextManager(
            engine=self,
            identity_name="__mock_engine",
            job_id=run_name,
            public_stickers={},
            private_stickers={},
        )

    def new_context(self):
        return self.fl_ctx_mgr.new_context()

    def fire_event(self, event_type: str, fl_ctx: FLContext):
        pass


CHECK_TEST_CASES = [
    ({"gpu": [1, 2, 3, 4]}, {"gpu": 1}, True, {"gpu": [1]}),
    ({"gpu": [1, 2, 3, 4]}, {"gpu": 4}, True, {"gpu": [1, 2, 3, 4]}),
    ({"gpu": [1]}, {"gpu": 1}, True, {"gpu": [1]}),
    ({"gpu": [1], "cpu": [1, 2, 3, 4, 5]}, {"gpu": 1, "cpu": 3}, True, {"gpu": [1], "cpu": [1, 2, 3]}),
    ({"gpu": [1]}, {"gpu": 2}, False, {}),
    ({"gpu": [1, 2]}, {"gpu": 5}, False, {}),
    ({"gpu": [1, 2]}, {"cpu": 1}, False, {}),
]


TEST_CASES = [
    ({"gpu": [1, 2, 3, 4]}, {"gpu": 1}, {"gpu": [1]}),
    ({"gpu": [1, 2, 3, 4]}, {"gpu": 4}, {"gpu": [1, 2, 3, 4]}),
    ({"gpu": [1]}, {"gpu": 1}, {"gpu": [1]}),
    ({"gpu": [1], "cpu": [1, 2, 3, 4, 5]}, {"gpu": 1, "cpu": 3}, {"gpu": [1], "cpu": [1, 2, 3]}),
]


class TestListResourceManager:
    @pytest.mark.parametrize(
        "resources, resource_requirement, expected_check_result, expected_reserved_resources", CHECK_TEST_CASES
    )
    def test_check_resource(self, resources, resource_requirement, expected_check_result, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx
            )
        assert expected_check_result == check_result
        if expected_check_result:
            assert expected_reserved_resources == list_resource_manager.reserved_resources[token][0]

    @pytest.mark.parametrize("resources, resource_requirement, expected_reserved_resources", TEST_CASES)
    def test_cancel_resource(self, resources, resource_requirement, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[token][0]
        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(
                resource_requirement=resource_requirement, token=token, fl_ctx=fl_ctx
            )
        assert list_resource_manager.reserved_resources == {}

    @pytest.mark.parametrize("resources, resource_requirement, expected_reserved_resources", TEST_CASES)
    def test_allocate_resource(self, resources, resource_requirement, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token, fl_ctx=fl_ctx
            )
        assert result == expected_reserved_resources

    @pytest.mark.parametrize("resources, resource_requirement, expected_reserved_resources", TEST_CASES)
    def test_free_resource(self, resources, resource_requirement, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx
            )
        assert expected_reserved_resources == list_resource_manager.reserved_resources[token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token, fl_ctx=fl_ctx
            )
        with engine.new_context() as fl_ctx:
            list_resource_manager.free_resources(resources=result, token=token, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}

    def test_check_one_check_two_then_allocate_two_allocate_one(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={"gpu": [f"gpu_{i}" for i in range(4)]})
        resource_requirement = {"gpu": 1}

        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx
            )
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx
            )

        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token2, fl_ctx=fl_ctx
            )
        assert result == {"gpu": ["gpu_1"]}

        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token1, fl_ctx=fl_ctx
            )
        assert result == {"gpu": ["gpu_0"]}

    def test_check_one_cancel_one_check_four_then_allocate_four(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={"gpu": [f"gpu_{i}" for i in range(4)]})
        resource_requirement1 = {"gpu": 1}
        resource_requirement2 = {"gpu": 4}

        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement1, fl_ctx=fl_ctx
            )

        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(
                resource_requirement=resource_requirement1, token=token1, fl_ctx=fl_ctx
            )

        with engine.new_context() as fl_ctx:
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement2, fl_ctx=fl_ctx
            )
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement2, token=token2, fl_ctx=fl_ctx
            )
        assert result == {"gpu": ["gpu_0", "gpu_1", "gpu_2", "gpu_3"]}

    def test_check_and_timeout(self):
        timeout = 5
        engine = MockEngine()
        list_resource_manager = ListResourceManager(
            resources={"gpu": [f"gpu_{i}" for i in range(4)]}, expiration_period=timeout
        )
        resource_requirement = {"gpu": 1}

        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.SYSTEM_START, fl_ctx=fl_ctx)
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx
            )
            assert {"gpu": ["gpu_0"]} == list_resource_manager.reserved_resources[token][0]
        time.sleep(timeout + 1)
        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.SYSTEM_END, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}
        assert list_resource_manager.resources == {"gpu": deque(["gpu_0", "gpu_1", "gpu_2", "gpu_3"])}

fcefb788c2f8da6c3d0a4a677a3d93b8967306e1
<mask token>


class MockEngine:
    <mask token>
    <mask token>

    def fire_event(self, event_type: str, fl_ctx: FLContext):
        pass


<mask token>


class TestListResourceManager:

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_check_result, expected_reserved_resources'
        , CHECK_TEST_CASES)
    def test_check_resource(self, resources, resource_requirement,
        expected_check_result, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_check_result == check_result
        if expected_check_result:
            assert expected_reserved_resources == list_resource_manager.reserved_resources[
                token][0]

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_cancel_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(resource_requirement=
                resource_requirement, token=token, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_allocate_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token,
                fl_ctx=fl_ctx)
        assert result == expected_reserved_resources

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_free_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token,
                fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            list_resource_manager.free_resources(resources=result, token=
                token, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}

    def test_check_one_check_two_then_allocate_two_allocate_one(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]})
        resource_requirement = {'gpu': 1}
        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token2,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_1']}
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token1,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_0']}

    def test_check_one_cancel_one_check_four_then_allocate_four(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]})
        resource_requirement1 = {'gpu': 1}
        resource_requirement2 = {'gpu': 4}
        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement1, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(resource_requirement=
                resource_requirement1, token=token1, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement2, fl_ctx=fl_ctx)
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement2, token=token2,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_0', 'gpu_1', 'gpu_2', 'gpu_3']}

    def test_check_and_timeout(self):
        timeout = 5
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]}, expiration_period=timeout)
        resource_requirement = {'gpu': 1}
        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.
                SYSTEM_START, fl_ctx=fl_ctx)
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
            assert {'gpu': ['gpu_0']
                } == list_resource_manager.reserved_resources[token][0]
        time.sleep(timeout + 1)
        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.
                SYSTEM_END, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}
        assert list_resource_manager.resources == {'gpu': deque(['gpu_0',
            'gpu_1', 'gpu_2', 'gpu_3'])}

<mask token>


class MockEngine:

    def __init__(self, run_name='exp1'):
        self.fl_ctx_mgr = FLContextManager(engine=self, identity_name=
            '__mock_engine', job_id=run_name, public_stickers={},
            private_stickers={})
    <mask token>

    def fire_event(self, event_type: str, fl_ctx: FLContext):
        pass


<mask token>


class TestListResourceManager:

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_check_result, expected_reserved_resources'
        , CHECK_TEST_CASES)
    def test_check_resource(self, resources, resource_requirement,
        expected_check_result, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_check_result == check_result
        if expected_check_result:
            assert expected_reserved_resources == list_resource_manager.reserved_resources[
                token][0]

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_cancel_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(resource_requirement=
                resource_requirement, token=token, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_allocate_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token,
                fl_ctx=fl_ctx)
        assert result == expected_reserved_resources

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_free_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token,
                fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            list_resource_manager.free_resources(resources=result, token=
                token, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}

    def test_check_one_check_two_then_allocate_two_allocate_one(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]})
        resource_requirement = {'gpu': 1}
        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token2,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_1']}
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token1,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_0']}

    def test_check_one_cancel_one_check_four_then_allocate_four(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]})
        resource_requirement1 = {'gpu': 1}
        resource_requirement2 = {'gpu': 4}
        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement1, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(resource_requirement=
                resource_requirement1, token=token1, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement2, fl_ctx=fl_ctx)
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement2, token=token2,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_0', 'gpu_1', 'gpu_2', 'gpu_3']}

    def test_check_and_timeout(self):
        timeout = 5
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]}, expiration_period=timeout)
        resource_requirement = {'gpu': 1}
        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.
                SYSTEM_START, fl_ctx=fl_ctx)
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
            assert {'gpu': ['gpu_0']
                } == list_resource_manager.reserved_resources[token][0]
        time.sleep(timeout + 1)
        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.
                SYSTEM_END, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}
        assert list_resource_manager.resources == {'gpu': deque(['gpu_0',
            'gpu_1', 'gpu_2', 'gpu_3'])}

<mask token>


class MockEngine:

    def __init__(self, run_name='exp1'):
        self.fl_ctx_mgr = FLContextManager(engine=self, identity_name=
            '__mock_engine', job_id=run_name, public_stickers={},
            private_stickers={})

    def new_context(self):
        return self.fl_ctx_mgr.new_context()

    def fire_event(self, event_type: str, fl_ctx: FLContext):
        pass


<mask token>


class TestListResourceManager:

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_check_result, expected_reserved_resources'
        , CHECK_TEST_CASES)
    def test_check_resource(self, resources, resource_requirement,
        expected_check_result, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_check_result == check_result
        if expected_check_result:
            assert expected_reserved_resources == list_resource_manager.reserved_resources[
                token][0]

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_cancel_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(resource_requirement=
                resource_requirement, token=token, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_allocate_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token,
                fl_ctx=fl_ctx)
        assert result == expected_reserved_resources

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_free_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token,
                fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            list_resource_manager.free_resources(resources=result, token=
                token, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}

    def test_check_one_check_two_then_allocate_two_allocate_one(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]})
        resource_requirement = {'gpu': 1}
        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token2,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_1']}
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token1,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_0']}

    def test_check_one_cancel_one_check_four_then_allocate_four(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]})
        resource_requirement1 = {'gpu': 1}
        resource_requirement2 = {'gpu': 4}
        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement1, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(resource_requirement=
                resource_requirement1, token=token1, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement2, fl_ctx=fl_ctx)
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement2, token=token2,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_0', 'gpu_1', 'gpu_2', 'gpu_3']}

    def test_check_and_timeout(self):
        timeout = 5
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]}, expiration_period=timeout)
        resource_requirement = {'gpu': 1}
        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.
                SYSTEM_START, fl_ctx=fl_ctx)
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
            assert {'gpu': ['gpu_0']
                } == list_resource_manager.reserved_resources[token][0]
        time.sleep(timeout + 1)
        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.
                SYSTEM_END, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}
        assert list_resource_manager.resources == {'gpu': deque(['gpu_0',
            'gpu_1', 'gpu_2', 'gpu_3'])}

<mask token>


class MockEngine:

    def __init__(self, run_name='exp1'):
        self.fl_ctx_mgr = FLContextManager(engine=self, identity_name=
            '__mock_engine', job_id=run_name, public_stickers={},
            private_stickers={})

    def new_context(self):
        return self.fl_ctx_mgr.new_context()

    def fire_event(self, event_type: str, fl_ctx: FLContext):
        pass


CHECK_TEST_CASES = [({'gpu': [1, 2, 3, 4]}, {'gpu': 1}, True, {'gpu': [1]}),
    ({'gpu': [1, 2, 3, 4]}, {'gpu': 4}, True, {'gpu': [1, 2, 3, 4]}), ({
    'gpu': [1]}, {'gpu': 1}, True, {'gpu': [1]}), ({'gpu': [1], 'cpu': [1, 
    2, 3, 4, 5]}, {'gpu': 1, 'cpu': 3}, True, {'gpu': [1], 'cpu': [1, 2, 3]
    }), ({'gpu': [1]}, {'gpu': 2}, False, {}), ({'gpu': [1, 2]}, {'gpu': 5},
    False, {}), ({'gpu': [1, 2]}, {'cpu': 1}, False, {})]
TEST_CASES = [({'gpu': [1, 2, 3, 4]}, {'gpu': 1}, {'gpu': [1]}), ({'gpu': [
    1, 2, 3, 4]}, {'gpu': 4}, {'gpu': [1, 2, 3, 4]}), ({'gpu': [1]}, {'gpu':
    1}, {'gpu': [1]}), ({'gpu': [1], 'cpu': [1, 2, 3, 4, 5]}, {'gpu': 1,
    'cpu': 3}, {'gpu': [1], 'cpu': [1, 2, 3]})]


class TestListResourceManager:

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_check_result, expected_reserved_resources'
        , CHECK_TEST_CASES)
    def test_check_resource(self, resources, resource_requirement,
        expected_check_result, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_check_result == check_result
        if expected_check_result:
            assert expected_reserved_resources == list_resource_manager.reserved_resources[
                token][0]

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_cancel_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(resource_requirement=
                resource_requirement, token=token, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_allocate_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token,
                fl_ctx=fl_ctx)
        assert result == expected_reserved_resources

    @pytest.mark.parametrize(
        'resources, resource_requirement, expected_reserved_resources',
        TEST_CASES)
    def test_free_resource(self, resources, resource_requirement,
        expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[
            token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token,
                fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            list_resource_manager.free_resources(resources=result, token=
                token, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}

    def test_check_one_check_two_then_allocate_two_allocate_one(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]})
        resource_requirement = {'gpu': 1}
        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token2,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_1']}
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token1,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_0']}

    def test_check_one_cancel_one_check_four_then_allocate_four(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]})
        resource_requirement1 = {'gpu': 1}
        resource_requirement2 = {'gpu': 4}
        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement1, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(resource_requirement=
                resource_requirement1, token=token1, fl_ctx=fl_ctx)
        with engine.new_context() as fl_ctx:
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement2, fl_ctx=fl_ctx)
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement2, token=token2,
                fl_ctx=fl_ctx)
        assert result == {'gpu': ['gpu_0', 'gpu_1', 'gpu_2', 'gpu_3']}

    def test_check_and_timeout(self):
        timeout = 5
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={'gpu': [
            f'gpu_{i}' for i in range(4)]}, expiration_period=timeout)
        resource_requirement = {'gpu': 1}
        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.
                SYSTEM_START, fl_ctx=fl_ctx)
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx)
            assert {'gpu': ['gpu_0']
                } == list_resource_manager.reserved_resources[token][0]
        time.sleep(timeout + 1)
        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.
                SYSTEM_END, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}
        assert list_resource_manager.resources == {'gpu': deque(['gpu_0',
            'gpu_1', 'gpu_2', 'gpu_3'])}

# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import time
from collections import deque

import pytest

from nvflare.apis.event_type import EventType
from nvflare.apis.fl_context import FLContext, FLContextManager
from nvflare.app_common.resource_managers.list_resource_manager import ListResourceManager


class MockEngine:
    def __init__(self, run_name="exp1"):
        self.fl_ctx_mgr = FLContextManager(
            engine=self,
            identity_name="__mock_engine",
            job_id=run_name,
            public_stickers={},
            private_stickers={},
        )

    def new_context(self):
        return self.fl_ctx_mgr.new_context()

    def fire_event(self, event_type: str, fl_ctx: FLContext):
        pass


CHECK_TEST_CASES = [
    ({"gpu": [1, 2, 3, 4]}, {"gpu": 1}, True, {"gpu": [1]}),
    ({"gpu": [1, 2, 3, 4]}, {"gpu": 4}, True, {"gpu": [1, 2, 3, 4]}),
    ({"gpu": [1]}, {"gpu": 1}, True, {"gpu": [1]}),
    ({"gpu": [1], "cpu": [1, 2, 3, 4, 5]}, {"gpu": 1, "cpu": 3}, True, {"gpu": [1], "cpu": [1, 2, 3]}),
    ({"gpu": [1]}, {"gpu": 2}, False, {}),
    ({"gpu": [1, 2]}, {"gpu": 5}, False, {}),
    ({"gpu": [1, 2]}, {"cpu": 1}, False, {}),
]


TEST_CASES = [
    ({"gpu": [1, 2, 3, 4]}, {"gpu": 1}, {"gpu": [1]}),
    ({"gpu": [1, 2, 3, 4]}, {"gpu": 4}, {"gpu": [1, 2, 3, 4]}),
    ({"gpu": [1]}, {"gpu": 1}, {"gpu": [1]}),
    ({"gpu": [1], "cpu": [1, 2, 3, 4, 5]}, {"gpu": 1, "cpu": 3}, {"gpu": [1], "cpu": [1, 2, 3]}),
]


class TestListResourceManager:
    @pytest.mark.parametrize(
        "resources, resource_requirement, expected_check_result, expected_reserved_resources", CHECK_TEST_CASES
    )
    def test_check_resource(self, resources, resource_requirement, expected_check_result, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx
            )
        assert expected_check_result == check_result
        if expected_check_result:
            assert expected_reserved_resources == list_resource_manager.reserved_resources[token][0]

    @pytest.mark.parametrize("resources, resource_requirement, expected_reserved_resources", TEST_CASES)
    def test_cancel_resource(self, resources, resource_requirement, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[token][0]
        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(
                resource_requirement=resource_requirement, token=token, fl_ctx=fl_ctx
            )
        assert list_resource_manager.reserved_resources == {}

    @pytest.mark.parametrize("resources, resource_requirement, expected_reserved_resources", TEST_CASES)
    def test_allocate_resource(self, resources, resource_requirement, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            _, token = list_resource_manager.check_resources(resource_requirement=resource_requirement, fl_ctx=fl_ctx)
        assert expected_reserved_resources == list_resource_manager.reserved_resources[token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token, fl_ctx=fl_ctx
            )
        assert result == expected_reserved_resources

    @pytest.mark.parametrize("resources, resource_requirement, expected_reserved_resources", TEST_CASES)
    def test_free_resource(self, resources, resource_requirement, expected_reserved_resources):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources=resources)
        with engine.new_context() as fl_ctx:
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx
            )
        assert expected_reserved_resources == list_resource_manager.reserved_resources[token][0]
        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token, fl_ctx=fl_ctx
            )
        with engine.new_context() as fl_ctx:
            list_resource_manager.free_resources(resources=result, token=token, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}

    def test_check_one_check_two_then_allocate_two_allocate_one(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={"gpu": [f"gpu_{i}" for i in range(4)]})
        resource_requirement = {"gpu": 1}

        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx
            )
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx
            )

        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token2, fl_ctx=fl_ctx
            )
        assert result == {"gpu": ["gpu_1"]}

        with engine.new_context() as fl_ctx:
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement, token=token1, fl_ctx=fl_ctx
            )
        assert result == {"gpu": ["gpu_0"]}

    def test_check_one_cancel_one_check_four_then_allocate_four(self):
        engine = MockEngine()
        list_resource_manager = ListResourceManager(resources={"gpu": [f"gpu_{i}" for i in range(4)]})
        resource_requirement1 = {"gpu": 1}
        resource_requirement2 = {"gpu": 4}

        with engine.new_context() as fl_ctx:
            check1, token1 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement1, fl_ctx=fl_ctx
            )

        with engine.new_context() as fl_ctx:
            list_resource_manager.cancel_resources(
                resource_requirement=resource_requirement1, token=token1, fl_ctx=fl_ctx
            )

        with engine.new_context() as fl_ctx:
            check2, token2 = list_resource_manager.check_resources(
                resource_requirement=resource_requirement2, fl_ctx=fl_ctx
            )
            result = list_resource_manager.allocate_resources(
                resource_requirement=resource_requirement2, token=token2, fl_ctx=fl_ctx
            )
        assert result == {"gpu": ["gpu_0", "gpu_1", "gpu_2", "gpu_3"]}

    def test_check_and_timeout(self):
        timeout = 5
        engine = MockEngine()
        list_resource_manager = ListResourceManager(
            resources={"gpu": [f"gpu_{i}" for i in range(4)]}, expiration_period=timeout
        )
        resource_requirement = {"gpu": 1}

        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.SYSTEM_START, fl_ctx=fl_ctx)
            check_result, token = list_resource_manager.check_resources(
                resource_requirement=resource_requirement, fl_ctx=fl_ctx
            )
            assert {"gpu": ["gpu_0"]} == list_resource_manager.reserved_resources[token][0]
        time.sleep(timeout + 1)
        with engine.new_context() as fl_ctx:
            list_resource_manager.handle_event(event_type=EventType.SYSTEM_END, fl_ctx=fl_ctx)
        assert list_resource_manager.reserved_resources == {}
        assert list_resource_manager.resources == {"gpu": deque(["gpu_0", "gpu_1", "gpu_2", "gpu_3"])}
