1# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
2#
3# This source code is licensed under the MIT license found in the
4# LICENSE file in the root directory of this source tree.
5
6import os
7import typing as tp
8import numpy as np
9from . import tools
10from . import testing
11
12
13@testing.parametrized(
14    void=([], []),
15    one=(["a"], []),
16    two=([1, 2], [(1, 2)]),
17    three=([1, 2, 3], [(1, 2), (2, 3)]),
18)
19def test_pairwise(iterator: tp.Iterable[tp.Any], expected: tp.List[tp.Tuple[tp.Any, ...]]) -> None:
20    output = list(tools.pairwise(iterator))
21    testing.printed_assert_equal(output, expected)
22
23
24def test_roundrobin() -> None:
25    output = list(tools.roundrobin([1, 2, 3], (x for x in [4, 5, 6, 7]), (8,)))
26    np.testing.assert_array_equal(output, [1, 4, 8, 2, 5, 3, 6, 7])
27
28
29def test_grouper() -> None:
30    output = list(tools.grouper("ABCDEFG", 3, "x"))
31    testing.printed_assert_equal(output, [list(x) for x in ["ABC", "DEF", "Gxx"]])
32
33
34def test_sleeper() -> None:
35    min_sleep = 1e-5
36    sleeper = tools.Sleeper(min_sleep=min_sleep)
37    np.testing.assert_almost_equal(sleeper._get_advised_sleep_duration(), min_sleep, decimal=5)
38    sleeper.start_timer()
39    # not precise enough to test for exact time
40    # np.testing.assert_almost_equal(sleeper._get_advised_sleep_duration(), min_sleep, decimal=5)
41    sleeper.stop_timer()
42    # np.testing.assert_almost_equal(sleeper._get_advised_sleep_duration(), min_sleep, decimal=5)
43
44
45def test_mutable_set() -> None:
46    s: tools.OrderedSet[int] = tools.OrderedSet((1, 2, 3))
47    assert tuple(s) == (1, 2, 3)
48    s.add(1)
49    s.add(4)
50    assert tuple(s) == (2, 3, 1, 4)
51    #
52    union = s & {3, 2}
53    assert isinstance(union, tools.OrderedSet)
54    assert tuple(union) == (2, 3)
55    #
56    assert s.pop() == 2
57    assert s.popright() == 4
58    assert tuple(s) == (3, 1)
59    #
60    s = tools.OrderedSet((1, 2, 3))
61    intersect = s | {5, 6}
62    assert isinstance(intersect, tools.OrderedSet)
63    assert tuple(intersect) == (1, 2, 3, 5, 6)
64    intersect = {5, 6} | s
65    assert isinstance(intersect, tools.OrderedSet)
66    assert tuple(intersect) == (1, 2, 3, 5, 6)  # same behavior, always appended to the end
67
68
69def test_set_env() -> None:
70    with tools.set_env(BLUBLU=1):
71        assert os.environ.get("BLUBLU", None) == "1"
72    assert os.environ.get("BLUBLU", None) is None
73
74
75def test_flatten() -> None:
76    out = tools.flatten(["a", {"truc": [4, 5]}])
77    assert out == {"0": "a", "1.truc.0": 4, "1.truc.1": 5}
78