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

1 import json

2 import logging

3 import os

4 from typing import Dict, Optional

5

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

7 from pydantic_settings import BaseSettings

8

9 logging.basicConfig(

10 level=logging.INFO,

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

12)

13 logger = logging.getLogger(__name__)

14

15

16 class 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

35 class 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

64 class 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

74 class 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

83 class 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

91 class 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()}")