Coverage for tests/test_OneRoom_CIA.py: 94%

1import pytest 

2import pandas as pd 

3import os 

4import sys 

5from pathlib import Path 

6import importlib.util 

7import json 

8from util import module_cleanup, round_floats_in_structure 

9 

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

11# This helps in locating the agentlib_flexquant package if needed 

12root_path = Path(__file__).parent.parent 

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

14 

15 

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

17 """Create 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 """ 

30 if df is None or df.empty: 

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

32 

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

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

35 

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

37 stats_dict_raw = summary_stats.to_dict() 

38 

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

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

41 stats_dict_clean = { 

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

43 for k, v in stats_dict_raw.items() 

44 } 

45 

46 # Create the final summary object 

47 summary = { 

48 "shape": df.shape, 

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

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

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

52 "statistics": stats_dict_clean, 

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

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

55 } 

56 return summary 

57 

58 

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

60 snapshot_name: str): 

61 """Assert that a DataFrame's summary matches a stored snapshot. 

62 

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

64 to compare it against a stored version. 

65 

66 """ 

67 # Create a summary of the dataframe 

68 summary = create_dataframe_summary(df) 

69 

70 # Round all numbers in the summary to handle cross-platform differences 

71 rounded_summary = round_floats_in_structure(summary, precision=1) 

72 

73 # Convert the summary dictionary to a formatted JSON string 

74 summary_json = json.dumps(rounded_summary, indent=2, sort_keys=True) 

75 

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

77 snapshot.assert_match(summary_json, snapshot_name) 

78 

79def run_example_from_path(example_path: Path): 

80 """Dynamically import and run the 'run_example' function from a script 

81 in the specified directory. 

82 

83 This function robustly handles changing the working directory AND the 

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

85 and its local modules. 

86 

87 """ 

88 run_script_path = example_path / 'main_cia_flex.py' 

89 if not run_script_path.is_file(): 

90 raise FileNotFoundError( 

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

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

93 ) 

94 

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

96 original_cwd = Path.cwd() 

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

98 

99 module_name = f"agentlib_flexquant.tests.examples.{example_path.name}" 

100 

101 try: 

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

103 os.chdir(example_path) 

104 

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

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

107 

108 # Dynamically import the run_example function from the script 

109 spec = importlib.util.spec_from_file_location(module_name, run_script_path) 

110 run_module = importlib.util.module_from_spec(spec) 

111 sys.modules[module_name] = run_module 

112 spec.loader.exec_module(run_module) 

113 

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

115 raise AttributeError( 

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

117 

118 # Execute the function and get the results 

119 results = run_module.run_example(until=3600) 

120 return results 

121 

122 finally: 

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

124 os.chdir(original_cwd) 

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

126 

127 

128def test_oneroom_cia(snapshot, module_cleanup): 

129 """Unit test for the OneRoom_CIA example using snapshot testing. 

130 

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

132 full resulting dataframes against stored snapshots. 

133 

134 """ 

135 # Define the path to the example directory 

136 example_path = root_path / 'examples' / 'OneRoom_CIA' 

137 

138 # Run the example and get the results object 

139 res = run_example_from_path(example_path) 

140 

141 # Extract the full resulting dataframes as requested 

142 df_neg_flex_res = res["NegFlexMPC"]["NegFlexMPC"] 

143 df_pos_flex_res = res["PosFlexMPC"]["PosFlexMPC"] 

144 df_baseline_res = res["myMPCAgent"]["Baseline"] 

145 df_indicator_res = res["FlexibilityIndicator"]["FlexibilityIndicator"] 

146 

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

148 assert_frame_matches_summary_snapshot( 

149 snapshot, 

150 df_neg_flex_res, 

151 'oneroom_cia_neg_flex_summary.json' 

152 ) 

153 assert_frame_matches_summary_snapshot( 

154 snapshot, 

155 df_pos_flex_res, 

156 'oneroom_cia_pos_flex_summary.json' 

157 ) 

158 assert_frame_matches_summary_snapshot( 

159 snapshot, 

160 df_baseline_res, 

161 'oneroom_cia_baseline_summary.json' 

162 ) 

163 assert_frame_matches_summary_snapshot( 

164 snapshot, 

165 df_indicator_res, 

166 'oneroom_cia_indicator_summary.json' 

167 )