Coverage for src/otg_mcp/config.py: 35%

135 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-05-19 00:42 -0700

1import json 

2import logging 

3import os 

4from typing import Dict, Optional 

5 

6from pydantic import BaseModel, ConfigDict, Field, validator, ValidationError 

7from pydantic_settings import BaseSettings 

8 

9logging.basicConfig( 

10 level=logging.INFO, 

11 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", 

12) 

13logger = logging.getLogger(__name__) 

14 

15 

16class LoggingConfig(BaseSettings): 

17 """Configuration for logging.""" 

18 

19 LOG_LEVEL: str = Field( 

20 default="INFO", 

21 description="Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)", 

22 ) 

23 

24 @validator("LOG_LEVEL") 

25 def validate_log_level(cls, v: str) -> str: 

26 valid_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] 

27 upper_v = v.upper() 

28 if upper_v not in valid_levels: 

29 logger.error(f"LOG_LEVEL must be one of {valid_levels}") 

30 raise ValueError(f"LOG_LEVEL must be one of {valid_levels}") 

31 logger.info(f"Validated log level: {upper_v}") 

32 return upper_v 

33 

34 

35class PortConfig(BaseModel): 

36 """Configuration for a port on a traffic generator.""" 

37 

38 location: Optional[str] = Field( 

39 None, description="Location of the port (hostname:port)" 

40 ) 

41 name: Optional[str] = Field(None, description="Name of the port") 

42 interface: Optional[str] = Field( 

43 None, description="Interface name (backward compatibility)" 

44 ) 

45 

46 @validator("location", pre=True, always=True) 

47 def validate_location(cls, v, values): 

48 """Validate location, using interface if location is not provided.""" 

49 if v is None and "interface" in values and values["interface"] is not None: 49 ↛ 50line 49 didn't jump to line 50 because the condition on line 49 was never true

50 return values["interface"] 

51 return v 

52 

53 @validator("name", pre=True, always=True) 

54 def validate_name(cls, v, values): 

55 """Validate name, using interface or location if name is not provided.""" 

56 if v is None: 56 ↛ 57line 56 didn't jump to line 57 because the condition on line 56 was never true

57 if "interface" in values and values["interface"] is not None: 

58 return values["interface"] 

59 if "location" in values and values["location"] is not None: 

60 return values["location"] 

61 return v 

62 

63 

64class TargetConfig(BaseModel): 

65 """Configuration for a traffic generator target.""" 

66 

67 ports: Dict[str, PortConfig] = Field( 

68 default_factory=dict, description="Port configurations mapped by port name" 

69 ) 

70 

71 model_config = ConfigDict(extra="forbid") 

72 

73 

74class TargetsConfig(BaseSettings): 

75 """Configuration for all available traffic generator targets.""" 

76 

77 targets: Dict[str, TargetConfig] = Field( 

78 default_factory=dict, 

79 description="Target configurations mapped by hostname:port", 

80 ) 

81 

82 

83class SchemaConfig(BaseSettings): 

84 """Configuration for schema handling.""" 

85 

86 schema_path: Optional[str] = Field( 

87 default=None, description="Path to directory containing custom schema files" 

88 ) 

89 

90 

91class Config: 

92 """Main configuration for the MCP server.""" 

93 

94 def __init__(self, config_file: Optional[str] = None): 

95 self.logging = LoggingConfig() 

96 self.targets = TargetsConfig() 

97 self.schemas = SchemaConfig() 

98 

99 logger.info("Initializing configuration") 

100 if config_file: 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true

101 logger.info(f"Loading configuration from file: {config_file}") 

102 self.load_config_file(config_file) 

103 elif not self.targets.targets: 103 ↛ exitline 103 didn't return from function '__init__' because the condition on line 103 was always true

104 logger.info("No targets defined - adding default development target") 

105 example_target = TargetConfig( 

106 ports={ 

107 "p1": PortConfig( 

108 location="localhost:5555", name="p1", interface=None 

109 ), 

110 "p2": PortConfig( 

111 location="localhost:5555", name="p2", interface=None 

112 ), 

113 } 

114 ) 

115 self.targets.targets["localhost:8443"] = example_target 

116 

117 def load_config_file(self, config_file_path: str) -> None: 

118 """ 

119 Load the traffic generator configuration from a JSON file. 

120 

121 Args: 

122 config_file_path: Path to the JSON configuration file 

123 

124 Raises: 

125 FileNotFoundError: If the config file doesn't exist 

126 json.JSONDecodeError: If the config file isn't valid JSON 

127 ValueError: If the config file doesn't have the expected structure 

128 """ 

129 logger.info(f"Loading traffic generator configuration from: {config_file_path}") 

130 

131 if not os.path.exists(config_file_path): 

132 error_msg = f"Configuration file not found: {config_file_path}" 

133 logger.critical(error_msg) 

134 raise FileNotFoundError(error_msg) 

135 

136 try: 

137 with open(config_file_path, "r") as file: 

138 config_data = json.load(file) 

139 

140 logger.info("Validating configuration structure") 

141 if "targets" not in config_data: 

142 error_msg = "Configuration file must contain a 'targets' property" 

143 logger.critical(error_msg) 

144 raise ValueError(error_msg) 

145 

146 logger.info("Clearing existing targets and initializing new configuration") 

147 self.targets = TargetsConfig() 

148 

149 logger.info("Processing each target in configuration") 

150 for hostname, target_data in config_data["targets"].items(): 

151 if not isinstance(target_data, dict) or "ports" not in target_data: 

152 error_msg = f"Target '{hostname}' must contain a 'ports' dictionary" 

153 logger.error(error_msg) 

154 continue 

155 

156 logger.info(f"Creating target config for {hostname}") 

157 

158 logger.info("Validating target configuration using Pydantic model") 

159 try: 

160 target_config = TargetConfig(**target_data) 

161 except ValidationError as e: 

162 error_msg = ( 

163 f"Invalid target configuration for '{hostname}': {str(e)}" 

164 ) 

165 logger.error(error_msg) 

166 if "extra fields not permitted" in str(e): 

167 logger.error( 

168 "The configuration contains fields that are not allowed. " 

169 "apiVersion should not be included in target configuration." 

170 ) 

171 continue 

172 

173 logger.info(f"Adding target {hostname} to configuration") 

174 self.targets.targets[hostname] = target_config 

175 

176 logger.info("Checking for schema path in configuration") 

177 if "schema_path" in config_data: 

178 schema_path = config_data["schema_path"] 

179 logger.info(f"Found schema_path in config: {schema_path}") 

180 if os.path.exists(schema_path): 

181 self.schemas.schema_path = schema_path 

182 logger.info(f"Using custom schema path: {schema_path}") 

183 else: 

184 logger.warning( 

185 f"Specified schema path does not exist: {schema_path}" 

186 ) 

187 

188 logger.info( 

189 f"Successfully loaded configuration with {len(self.targets.targets)} targets" 

190 ) 

191 

192 except json.JSONDecodeError as e: 

193 error_msg = f"Invalid JSON in configuration file: {str(e)}" 

194 logger.critical(error_msg) 

195 raise 

196 except Exception as e: 

197 error_msg = f"Error loading configuration: {str(e)}" 

198 logger.critical(error_msg) 

199 raise 

200 

201 def setup_logging(self): 

202 """Configure logging based on the provided settings.""" 

203 try: 

204 log_level = getattr(logging, self.logging.LOG_LEVEL) 

205 print(f"Setting up logging at level {self.logging.LOG_LEVEL}") 

206 

207 logger.info( 

208 "Setting up both basic config and console handler for comprehensive logging" 

209 ) 

210 logging.basicConfig( 

211 level=log_level, 

212 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", 

213 ) 

214 

215 logger.info("Configuring root logger") 

216 root_logger = logging.getLogger() 

217 root_logger.setLevel(log_level) 

218 

219 logger.info(f"Setting module logger to level {log_level}") 

220 module_logger = logging.getLogger("otg_mcp") 

221 module_logger.setLevel(log_level) 

222 

223 logger.info("Checking if root logger has handlers, adding if needed") 

224 if not root_logger.handlers: 

225 console_handler = logging.StreamHandler() 

226 console_handler.setLevel(log_level) 

227 formatter = logging.Formatter( 

228 "%(asctime)s - %(name)s - %(levelname)s - %(message)s" 

229 ) 

230 console_handler.setFormatter(formatter) 

231 root_logger.addHandler(console_handler) 

232 print("Added console handler to root logger") 

233 

234 logger.info("Logging system initialized with handlers and formatters") 

235 logger.info(f"Logging configured at level {self.logging.LOG_LEVEL}") 

236 except Exception as e: 

237 print(f"CRITICAL ERROR setting up logging: {str(e)}") 

238 import traceback 

239 

240 print(f"Stack trace: {traceback.format_exc()}") 

241 logger.critical(f"Failed to set up logging: {str(e)}") 

242 logger.critical(f"Stack trace: {traceback.format_exc()}")