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

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

2

3 from datetime import datetime

4 from typing import Any, Dict, Optional

5

6 from pydantic import BaseModel, Field

7

8

9 class ApiResponse(BaseModel):

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

11

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

13

14

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

23 class MetricsResponse(ApiResponse):

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

25

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

27

28

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

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

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

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

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

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

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

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

116 class SnappiError(BaseModel):

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

118

119 error: str

120 detail: Optional[str] = None

121 code: Optional[int] = None