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"""Unit tests for passing shell environment variables into LocalEnv scripts.""" 

6import pytest 

7 

8from mlos_bench.environments.local.local_fileshare_env import LocalFileShareEnv 

9from mlos_bench.services.config_persistence import ConfigPersistenceService 

10from mlos_bench.services.local.local_exec import LocalExecService 

11from mlos_bench.tests.services.remote.mock.mock_fileshare_service import ( 

12 MockFileShareService, 

13) 

14from mlos_bench.tunables.tunable_groups import TunableGroups 

15 

16# pylint: disable=redefined-outer-name 

17 

18 

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

20def mock_fileshare_service() -> MockFileShareService: 

21 """Create a new mock FileShareService instance.""" 

22 return MockFileShareService( 

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

24 parent=LocalExecService(parent=ConfigPersistenceService()), 

25 ) 

26 

27 

28@pytest.fixture 

29def local_fileshare_env( 

30 tunable_groups: TunableGroups, 

31 mock_fileshare_service: MockFileShareService, 

32) -> LocalFileShareEnv: 

33 """Create a LocalFileShareEnv instance.""" 

34 env = LocalFileShareEnv( 

35 name="TestLocalFileShareEnv", 

36 config={ 

37 "const_args": { 

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

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

40 }, 

41 "tunable_params": ["boot"], 

42 "shell_env_params": [ 

43 "trial_id", # From "const_arg" 

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

45 ], 

46 "upload": [ 

47 { 

48 "from": "grub.cfg", 

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

50 }, 

51 { 

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

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

54 }, 

55 ], 

56 "run": ["echo No-op run"], 

57 "download": [ 

58 { 

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

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

61 }, 

62 ], 

63 }, 

64 tunables=tunable_groups, 

65 service=mock_fileshare_service, 

66 ) 

67 return env 

68 

69 

70def test_local_fileshare_env( 

71 tunable_groups: TunableGroups, 

72 mock_fileshare_service: MockFileShareService, 

73 local_fileshare_env: LocalFileShareEnv, 

74) -> None: 

75 """Test that the LocalFileShareEnv correctly expands the `$VAR` variables in the 

76 upload and download sections of the config. 

77 """ 

78 with local_fileshare_env as env_context: 

79 assert env_context.setup(tunable_groups) 

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

81 assert status.is_succeeded() 

82 assert mock_fileshare_service.get_upload() == [ 

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

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

85 ] 

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

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

88 assert mock_fileshare_service.get_download() == [ 

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

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

91 ]