Coverage for mlos_bench/mlos_bench/tests/environments/__init__.py: 100%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Tests helpers for mlos_bench.environments. 

7""" 

8 

9from datetime import datetime 

10from typing import Any, Dict, List, Optional, Tuple 

11 

12import pytest 

13 

14from mlos_bench.environments.base_environment import Environment 

15 

16from mlos_bench.tunables.tunable import TunableValue 

17from mlos_bench.tunables.tunable_groups import TunableGroups 

18 

19 

20def check_env_success(env: Environment, 

21 tunable_groups: TunableGroups, 

22 expected_results: Dict[str, TunableValue], 

23 expected_telemetry: List[Tuple[datetime, str, Any]], 

24 global_config: Optional[dict] = None) -> None: 

25 """ 

26 Set up an environment and run a test experiment there. 

27 

28 Parameters 

29 ---------- 

30 tunable_groups : TunableGroups 

31 Tunable parameters (usually come from a fixture). 

32 env : Environment 

33 An environment to query for the results. 

34 expected_results : Dict[str, float] 

35 Expected results of the benchmark. 

36 expected_telemetry : List[Tuple[datetime, str, Any]] 

37 Expected telemetry data of the benchmark. 

38 global_config : dict 

39 Global params. 

40 """ 

41 with env as env_context: 

42 

43 assert env_context.setup(tunable_groups, global_config) 

44 

45 (status, _ts, data) = env_context.run() 

46 assert status.is_succeeded() 

47 assert data == pytest.approx(expected_results, nan_ok=True) 

48 

49 (status, _ts, telemetry) = env_context.status() 

50 assert status.is_good() 

51 assert telemetry == pytest.approx(expected_telemetry, nan_ok=True) 

52 

53 env_context.teardown() 

54 assert not env_context._is_ready # pylint: disable=protected-access 

55 

56 

57def check_env_fail_telemetry(env: Environment, tunable_groups: TunableGroups) -> None: 

58 """ 

59 Set up a local environment and run a test experiment there; 

60 Make sure the environment `.status()` call fails. 

61 

62 Parameters 

63 ---------- 

64 tunable_groups : TunableGroups 

65 Tunable parameters (usually come from a fixture). 

66 env : Environment 

67 An environment to query for the results. 

68 """ 

69 with env as env_context: 

70 

71 assert env_context.setup(tunable_groups) 

72 (status, _ts, _data) = env_context.run() 

73 assert status.is_succeeded() 

74 

75 with pytest.raises(ValueError): 

76 env_context.status()