Coverage for src/otg_mcp/models/models.py: 98%

48 statements  

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

1"""Unified models for OTG MCP.""" 

2 

3from datetime import datetime 

4from typing import Any, Dict, Optional 

5 

6from pydantic import BaseModel, Field 

7 

8 

9class ApiResponse(BaseModel): 

10 """Base response model for API responses.""" 

11 

12 status: str = Field(default="success", description="Status of the response") 

13 

14 

15class ConfigResponse(ApiResponse): 

16 """Response model for configuration operations.""" 

17 

18 config: Optional[Dict[str, Any]] = Field( 

19 default=None, description="Configuration data" 

20 ) 

21 

22 

23class MetricsResponse(ApiResponse): 

24 """Response model for metrics operations.""" 

25 

26 metrics: Optional[Dict[str, Any]] = Field(default=None, description="Metrics data") 

27 

28 

29class CaptureResponse(ApiResponse): 

30 """Response model for capture operations.""" 

31 

32 port: str = Field(..., description="Name of the port used for capture") 

33 data: Optional[Dict[str, Any]] = Field(default=None, description="Capture data") 

34 capture_id: Optional[str] = None 

35 file_path: Optional[str] = Field( 

36 default=None, description="Path to the saved capture file (.pcap)" 

37 ) 

38 

39 

40class ControlResponse(ApiResponse): 

41 """Response model for control responses.""" 

42 

43 action: str = Field(..., description="Action that was performed") 

44 verified: Optional[bool] = Field( 

45 default=None, description="Whether the action was verified" 

46 ) 

47 timestamp: datetime = Field(default_factory=datetime.now) 

48 result: Optional[Dict[str, Any]] = Field( 

49 default=None, description="Result details of the control operation" 

50 ) 

51 

52 

53class PortInfo(BaseModel): 

54 """Information about a port on a traffic generator.""" 

55 

56 name: str = Field(..., description="Name of the port") 

57 location: str = Field(..., description="Location of the port (hostname:port)") 

58 interface: Optional[str] = Field( 

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

60 ) 

61 

62 @property 

63 def interface_name(self) -> str: 

64 """Get the interface name, falling back to location if not set.""" 

65 return self.interface or self.location 

66 

67 

68class CapabilitiesVersionResponse(BaseModel): 

69 """Response from the capabilities/version endpoint.""" 

70 

71 api_spec_version: str 

72 sdk_version: str 

73 app_version: str 

74 

75 

76class TrafficGeneratorInfo(BaseModel): 

77 """Information about a traffic generator.""" 

78 

79 hostname: str = Field(..., description="Hostname of the traffic generator") 

80 ports: Dict[str, PortInfo] = Field( 

81 default_factory=dict, description="Ports available on this generator" 

82 ) 

83 available: bool = Field( 

84 default=True, description="Whether the generator is available" 

85 ) 

86 

87 

88class TrafficGeneratorStatus(ApiResponse): 

89 """Status of all traffic generators.""" 

90 

91 generators: Dict[str, TrafficGeneratorInfo] = Field( 

92 default_factory=dict, description="All available traffic generators" 

93 ) 

94 

95 

96class TargetHealthInfo(BaseModel): 

97 """Health information for a traffic generator target.""" 

98 

99 name: str = Field(..., description="Name of the target") 

100 healthy: bool = Field(..., description="Whether the target is healthy") 

101 version_info: Optional[CapabilitiesVersionResponse] = Field( 

102 None, description="Version information when available" 

103 ) 

104 error: Optional[str] = Field(None, description="Error message if unhealthy") 

105 

106 

107class HealthStatus(BaseModel): 

108 """Health status collection of all traffic generators.""" 

109 

110 status: str = Field(default="success", description="Status of the response") 

111 targets: Dict[str, TargetHealthInfo] = Field( 

112 default_factory=dict, description="Health status of individual targets" 

113 ) 

114 

115 

116class SnappiError(BaseModel): 

117 """Error model for snappi errors.""" 

118 

119 error: str 

120 detail: Optional[str] = None 

121 code: Optional[int] = None