Coverage for tests/test_oneRoom_SimpleMPC.py: 94%

1import pytest 

2import pandas as pd 

3import os 

4import sys 

5from pathlib import Path 

6import importlib.util 

7import json 

8 

9# Add the project root to the Python path to allow for absolute imports 

10# This helps in locating the agentlib_flexquant package if needed 

11root_path = Path(__file__).parent.parent 

12sys.path.insert(0, str(root_path)) 

13 

14 

15def create_dataframe_summary(df: pd.DataFrame, precision: int = 6) -> dict: 

16 """ 

17 Creates a robust, compact summary of a DataFrame for snapshotting. 

18 

19 This summary is designed to be insensitive to minor floating-point differences 

20 while being highly sensitive to meaningful data changes. 

21 

22 Args: 

23 df: The pandas DataFrame to summarize. 

24 precision: The number of decimal places to round float values to. 

25 

26 Returns: 

27 A dictionary containing the summary. 

28 """ 

29 if df is None or df.empty: 

30 return {"error": "DataFrame is empty or None"} 

31 

32 # Get descriptive statistics and round them to handle float precision issues 

33 summary_stats = df.describe().round(precision) 

34 

35 # Convert the stats DataFrame to a dictionary. This may have tuple keys. 

36 stats_dict_raw = summary_stats.to_dict() 

37 

38 # Create a new dictionary, converting any tuple keys into strings. 

39 # e.g., ('lower', 'P_el') becomes 'lower.P_el' 

40 stats_dict_clean = { 

41 ".".join(map(str, k)) if isinstance(k, tuple) else str(k): v 

42 for k, v in stats_dict_raw.items() 

43 } 

44 

45 # Create the final summary object 

46 summary = { 

47 "shape": df.shape, 

48 "columns": df.columns.tolist(), 

49 "index_start": str(df.index.min()), 

50 "index_end": str(df.index.max()), 

51 "statistics": stats_dict_clean, 

52 "head_5_rows": df.head(5).round(precision).to_dict(orient='split'), 

53 "tail_5_rows": df.tail(5).round(precision).to_dict(orient='split'), 

54 } 

55 return summary 

56 

57def assert_frame_matches_summary_snapshot(snapshot, df: pd.DataFrame, 

58 snapshot_name: str): 

59 """ 

60 Asserts that a DataFrame's summary matches a stored snapshot. 

61 

62 This function creates a summary of the dataframe and uses pytest-snapshot 

63 to compare it against a stored version. 

64 """ 

65 # Create a summary of the dataframe 

66 summary = create_dataframe_summary(df) 

67 

68 # Convert the summary dictionary to a formatted JSON string 

69 summary_json = json.dumps(summary, indent=2, sort_keys=True) 

70 

71 # Use snapshot.assert_match on the small, stable JSON string 

72 snapshot.assert_match(summary_json, snapshot_name) 

73 

74def run_example_from_path(example_path: Path): 

75 """ 

76 Dynamically imports and runs the 'run_example' function from a script 

77 in the specified directory. 

78 

79 This function robustly handles changing the working directory AND the 

80 Python import path, ensuring the script can find both its local files 

81 and its local modules. 

82 """ 

83 run_script_path = example_path / 'main_one_room_flex.py' 

84 if not run_script_path.is_file(): 

85 raise FileNotFoundError( 

86 f"Could not find the run script at {run_script_path}. " 

87 "Please ensure it is named 'run.py' or adjust the test code." 

88 ) 

89 

90 # --- SETUP: Store original paths before changing them --- 

91 original_cwd = Path.cwd() 

92 original_sys_path = sys.path[:] # Create a copy of the sys.path list 

93 

94 try: 

95 # --- STEP 1: Change CWD for file access (e.g., config.json) --- 

96 os.chdir(example_path) 

97 

98 # --- STEP 2: Add example dir to sys.path for module imports --- 

99 sys.path.insert(0, str(example_path)) 

100 

101 # Dynamically import the run_example function from the script 

102 spec = importlib.util.spec_from_file_location("run_module", run_script_path) 

103 run_module = importlib.util.module_from_spec(spec) 

104 spec.loader.exec_module(run_module) 

105 

106 if not hasattr(run_module, 'run_example'): 

107 raise AttributeError( 

108 "The 'run.py' script must contain a 'run_example' function.") 

109 

110 # Execute the function and get the results 

111 results = run_module.run_example(until=3600) 

112 return results 

113 

114 finally: 

115 # --- TEARDOWN: Always restore original paths to avoid side-effects --- 

116 os.chdir(original_cwd) 

117 sys.path[:] = original_sys_path # Restore the original sys.path 

118 

119 

120def test_oneroom_simple_mpc(snapshot): 

121 """ 

122 Unit test for the oneroom_simpleMPC example using snapshot testing. 

123 

124 This test runs the example via its own run script and compares the 

125 full resulting dataframes against stored snapshots. 

126 """ 

127 # Define the path to the example directory 

128 example_path = root_path / 'examples' / 'OneRoom_SimpleMPC' 

129 

130 # Run the example and get the results object 

131 res = run_example_from_path(example_path) 

132 

133 # Extract the full resulting dataframes as requested 

134 df_neg_flex_res = res["NegFlexMPC"]["NegFlexMPC"] 

135 df_pos_flex_res = res["PosFlexMPC"]["PosFlexMPC"] 

136 df_baseline_res = res["FlexModel"]["Baseline"] 

137 df_indicator_res = res["FlexibilityIndicator"]["FlexibilityIndicator"] 

138 

139 # Assert that a summary of each result DataFrame matches its snapshot 

140 assert_frame_matches_summary_snapshot( 

141 snapshot, 

142 df_neg_flex_res, 

143 'oneroom_simpleMPC_neg_flex_summary.json' 

144 ) 

145 assert_frame_matches_summary_snapshot( 

146 snapshot, 

147 df_pos_flex_res, 

148 'oneroom_simpleMPC_pos_flex_summary.json' 

149 ) 

150 assert_frame_matches_summary_snapshot( 

151 snapshot, 

152 df_baseline_res, 

153 'oneroom_simpleMPC_baseline_summary.json' 

154 ) 

155 assert_frame_matches_summary_snapshot( 

156 snapshot, 

157 df_indicator_res, 

158 'oneroom_simpleMPC_indicator_summary.json' 

159 )