Coverage for mlos_bench/mlos_bench/tests/environments/local/local_fileshare_env_test.py: 100%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Unit tests for passing shell environment variables into LocalEnv scripts. 

7""" 

8import pytest 

9 

10from mlos_bench.tunables.tunable_groups import TunableGroups 

11from mlos_bench.environments.local.local_fileshare_env import LocalFileShareEnv 

12from mlos_bench.services.config_persistence import ConfigPersistenceService 

13from mlos_bench.services.local.local_exec import LocalExecService 

14 

15from mlos_bench.tests.services.remote.mock.mock_fileshare_service import MockFileShareService 

16 

17# pylint: disable=redefined-outer-name 

18 

19 

20@pytest.fixture(scope="module") 

21def mock_fileshare_service() -> MockFileShareService: 

22 """ 

23 Create a new mock FileShareService instance. 

24 """ 

25 return MockFileShareService( 

26 config={"fileShareName": "MOCK_FILESHARE"}, 

27 parent=LocalExecService(parent=ConfigPersistenceService()) 

28 ) 

29 

30 

31@pytest.fixture 

32def local_fileshare_env(tunable_groups: TunableGroups, 

33 mock_fileshare_service: MockFileShareService) -> LocalFileShareEnv: 

34 """ 

35 Create a LocalFileShareEnv instance. 

36 """ 

37 env = LocalFileShareEnv( 

38 name="TestLocalFileShareEnv", 

39 config={ 

40 "const_args": { 

41 "experiment_id": "EXP_ID", # Passed into "shell_env_params" 

42 "trial_id": 222, # NOT passed into "shell_env_params" 

43 }, 

44 "tunable_params": ["boot"], 

45 "shell_env_params": [ 

46 "trial_id", # From "const_arg" 

47 "idle", # From "tunable_params", == "halt" 

48 ], 

49 "upload": [ 

50 { 

51 "from": "grub.cfg", 

52 "to": "$experiment_id/$trial_id/input/grub.cfg", 

53 }, 

54 { 

55 "from": "data_$idle.csv", 

56 "to": "$experiment_id/$trial_id/input/data_$idle.csv", 

57 }, 

58 ], 

59 "run": [ 

60 "echo No-op run" 

61 ], 

62 "download": [ 

63 { 

64 "from": "$experiment_id/$trial_id/$idle/data.csv", 

65 "to": "output/data_$idle.csv", 

66 }, 

67 ], 

68 }, 

69 tunables=tunable_groups, 

70 service=mock_fileshare_service, 

71 ) 

72 return env 

73 

74 

75def test_local_fileshare_env(tunable_groups: TunableGroups, 

76 mock_fileshare_service: MockFileShareService, 

77 local_fileshare_env: LocalFileShareEnv) -> None: 

78 """ 

79 Test that the LocalFileShareEnv correctly expands the `$VAR` variables 

80 in the upload and download sections of the config. 

81 """ 

82 with local_fileshare_env as env_context: 

83 assert env_context.setup(tunable_groups) 

84 (status, _ts, _output) = env_context.run() 

85 assert status.is_succeeded() 

86 assert mock_fileshare_service.get_upload() == [ 

87 ("grub.cfg", "EXP_ID/222/input/grub.cfg"), 

88 ("data_halt.csv", "EXP_ID/222/input/data_halt.csv"), 

89 ] 

90 # NOTE: The "download" section is run twice -- once to check 

91 # the status of the run, and once to get the final results. 

92 assert mock_fileshare_service.get_download() == [ 

93 ("EXP_ID/222/halt/data.csv", "output/data_halt.csv"), 

94 ("EXP_ID/222/halt/data.csv", "output/data_halt.csv"), 

95 ]