WaterBenchmarkHub/tests/test_leakg3pd.py at main · HammerLabML/WaterBenchmarkHub · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
"""
Module provides tests to test the `water_benchmark_hub.leakg3pd.leakg3pd` module.
"""
from epyt_flow.simulation import ScenarioSimulator
from water_benchmark_hub import load

from .utils import get_temp_folder


def test_leakg3pd_net1():
    res = load("LeakG3PD")

    assert res.get_meta_info() != {}

    configs = res.load_scenarios(download_dir=get_temp_folder(), scenarios_id=range(1, 5),
                                 network="net1")
    for c in configs:
        with ScenarioSimulator(scenario_config=c) as sim:
            scada_data = sim.run_simulation()
            assert scada_data is not None

    data = res.load_data(download_dir=get_temp_folder(), scenarios_id=range(1, 5), network="net1")
    assert data is not None

    data = res.load_data(download_dir=get_temp_folder(), scenarios_id=range(1, 5),
                         network="net1", return_X_y=True,
                         return_features_desc=True, return_leak_locations=True)
    assert data is not None


def test_leakg3pd_net3():
    res = load("LeakG3PD")

    assert res.get_meta_info() != {}

    configs = res.load_scenarios(download_dir=get_temp_folder(), scenarios_id=range(1, 5),
                                 network="net3")
    for c in configs:
        with ScenarioSimulator(scenario_config=c) as sim:
            scada_data = sim.run_simulation()
            assert scada_data is not None

    data = res.load_data(download_dir=get_temp_folder(), scenarios_id=range(1, 5), network="net3")
    assert data is not None

    data = res.load_data(download_dir=get_temp_folder(), scenarios_id=range(1, 5),
                         network="net3", return_X_y=True,
                         return_features_desc=True, return_leak_locations=True)
    assert data is not None


def test_leakg3pd_hanoi():
    res = load("LeakG3PD")

    assert res.get_meta_info() != {}

    configs = res.load_scenarios(download_dir=get_temp_folder(), scenarios_id=range(1, 5),
                                 network="hanoi")
    for c in configs:
        with ScenarioSimulator(scenario_config=c) as sim:
            scada_data = sim.run_simulation()
            assert scada_data is not None

    data = res.load_data(download_dir=get_temp_folder(), scenarios_id=range(1, 5), network="hanoi")
    assert data is not None

    data = res.load_data(download_dir=get_temp_folder(), scenarios_id=range(1, 5),
                         network="hanoi", return_X_y=True,
                         return_features_desc=True, return_leak_locations=True)
    assert data is not None