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"""Tests helpers for mlos_bench.environments.""" 

6 

7from datetime import datetime 

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

9 

10import pytest 

11 

12from mlos_bench.environments.base_environment import Environment 

13from mlos_bench.tunables.tunable import TunableValue 

14from mlos_bench.tunables.tunable_groups import TunableGroups 

15 

16 

17def check_env_success( 

18 env: Environment, 

19 tunable_groups: TunableGroups, 

20 expected_results: Dict[str, TunableValue], 

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

22 global_config: Optional[dict] = None, 

23) -> None: 

24 """ 

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

26 

27 Parameters 

28 ---------- 

29 tunable_groups : TunableGroups 

30 Tunable parameters (usually come from a fixture). 

31 env : Environment 

32 An environment to query for the results. 

33 expected_results : Dict[str, float] 

34 Expected results of the benchmark. 

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

36 Expected telemetry data of the benchmark. 

37 global_config : dict 

38 Global params. 

39 """ 

40 with env as env_context: 

41 

42 assert env_context.setup(tunable_groups, global_config) 

43 

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

45 assert status.is_succeeded() 

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

47 

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

49 assert status.is_good() 

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

51 

52 env_context.teardown() 

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

54 

55 

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

57 """ 

58 Set up a local environment and run a test experiment there; Make sure the 

59 environment `.status()` call fails. 

60 

61 Parameters 

62 ---------- 

63 tunable_groups : TunableGroups 

64 Tunable parameters (usually come from a fixture). 

65 env : Environment 

66 An environment to query for the results. 

67 """ 

68 with env as env_context: 

69 

70 assert env_context.setup(tunable_groups) 

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

72 assert status.is_succeeded() 

73 

74 with pytest.raises(ValueError): 

75 env_context.status()