Coverage for src/usaspending/queries/sub_agency_query.py: 100%

1"""Sub-agency query implementation for retrieving agency sub-agencies with pagination.""" 

2 

3from typing import TYPE_CHECKING, Optional, Dict, Any, List 

4from .single_resource_base import SingleResourceBase 

5from ..exceptions import ValidationError 

6from ..client import USASpending 

7from ..logging_config import USASpendingLogger 

8 

9if TYPE_CHECKING: 

10 pass 

11 

12logger = USASpendingLogger.get_logger(__name__) 

13 

14 

15class SubAgencyQuery(SingleResourceBase): 

16 """Retrieve sub-agency data from the USAspending API. 

17  

18 This query class handles fetching sub-agency information including 

19 transaction counts and obligations filtered by fiscal year, agency type, 

20 and award type codes. 

21 """ 

22 

23 def __init__(self, client: USASpending): 

24 """Initialize SubAgencyQuery with client. 

25  

26 Args: 

27 client: USASpending client instance 

28 """ 

29 super().__init__(client) 

30 logger.debug("SubAgencyQuery initialized with client: %s", client) 

31 

32 @property 

33 def _endpoint(self) -> str: 

34 """Base endpoint for sub-agency retrieval.""" 

35 return '/agency/{toptier_code}/sub_agency/' 

36 

37 def _construct_endpoint(self, resource_id: str) -> str: 

38 """Construct the full endpoint URL for agency sub-agency. 

39  

40 Args: 

41 resource_id: The toptier_code for the agency 

42  

43 Returns: 

44 Full endpoint path including /sub_agency/ 

45 """ 

46 endpoint = self._endpoint.replace("{toptier_code}", resource_id) 

47 return endpoint 

48 

49 def find_by_id(self, toptier_code: str) -> Dict[str, Any]: 

50 """Not used for sub-agency - use get_subagencies instead. 

51  

52 Raises: 

53 NotImplementedError: This method should not be used directly 

54 """ 

55 raise NotImplementedError( 

56 "Use get_subagencies() method instead for sub-agency data" 

57 ) 

58 

59 def get_subagencies( 

60 self, 

61 toptier_code: str, 

62 fiscal_year: Optional[int] = None, 

63 agency_type: str = "awarding", 

64 award_type_codes: Optional[List[str]] = None, 

65 order: str = "desc", 

66 sort: str = "total_obligations", 

67 page: int = 1, 

68 limit: int = 100 

69 ) -> Dict[str, Any]: 

70 """Retrieve agency sub-agencies with optional filters and pagination. 

71  

72 Args: 

73 toptier_code: The toptier code of an agency (3-4 digit string) 

74 fiscal_year: Optional fiscal year for the data (defaults to current) 

75 agency_type: "awarding" or "funding" (defaults to "awarding") 

76 award_type_codes: Optional list of award type codes to filter by 

77 order: Sort direction "asc" or "desc" (defaults to "desc") 

78 sort: Sort field - "name", "total_obligations", "transaction_count", "new_award_count" 

79 page: Page number (defaults to 1) 

80 limit: Items per page, max 100 (defaults to 100) 

81  

82 Returns: 

83 Dictionary containing: 

84 - toptier_code: Agency toptier code 

85 - fiscal_year: Fiscal year of data 

86 - page_metadata: Pagination metadata 

87 - results: List of sub-agency data 

88 - messages: Any API messages 

89  

90 Raises: 

91 ValidationError: If toptier_code is invalid or other parameters are invalid 

92 APIError: If API request fails 

93 """ 

94 # Validate toptier_code 

95 if not toptier_code: 

96 raise ValidationError("toptier_code is required") 

97 

98 toptier_code = str(toptier_code).strip() 

99 if not toptier_code.isdigit() or len(toptier_code) not in [3, 4]: 

100 raise ValidationError( 

101 f"Invalid toptier_code: {toptier_code}. " 

102 "Must be a 3-4 digit numeric string" 

103 ) 

104 

105 # Validate agency_type 

106 if agency_type not in ["awarding", "funding"]: 

107 raise ValidationError( 

108 f"Invalid agency_type: {agency_type}. " 

109 "Must be 'awarding' or 'funding'" 

110 ) 

111 

112 # Validate order 

113 if order not in ["asc", "desc"]: 

114 raise ValidationError( 

115 f"Invalid order: {order}. " 

116 "Must be 'asc' or 'desc'" 

117 ) 

118 

119 # Validate sort 

120 valid_sorts = ["name", "total_obligations", "transaction_count", "new_award_count"] 

121 if sort not in valid_sorts: 

122 raise ValidationError( 

123 f"Invalid sort: {sort}. " 

124 f"Must be one of: {', '.join(valid_sorts)}" 

125 ) 

126 

127 # Validate page 

128 if page < 1: 

129 raise ValidationError("page must be >= 1") 

130 

131 # Validate limit 

132 if limit < 1 or limit > 100: 

133 raise ValidationError("limit must be between 1 and 100") 

134 

135 logger.debug( 

136 "Fetching sub-agencies for toptier_code: %s, fiscal_year: %s, " 

137 "agency_type: %s, award_type_codes: %s, order: %s, sort: %s, " 

138 "page: %s, limit: %s", 

139 toptier_code, fiscal_year, agency_type, award_type_codes, 

140 order, sort, page, limit 

141 ) 

142 

143 # Build params 

144 params = { 

145 "agency_type": agency_type, 

146 "order": order, 

147 "sort": sort, 

148 "page": page, 

149 "limit": limit 

150 } 

151 

152 if fiscal_year is not None: 

153 params["fiscal_year"] = fiscal_year 

154 

155 if award_type_codes: 

156 # Convert to list if needed and filter out None/empty values 

157 if isinstance(award_type_codes, (set, frozenset)): 

158 award_type_codes = list(award_type_codes) 

159 award_type_codes = [code for code in award_type_codes if code] 

160 

161 if award_type_codes: 

162 # API expects award_type_codes as array parameter 

163 params["award_type_codes"] = award_type_codes 

164 

165 # Construct endpoint 

166 endpoint = self._construct_endpoint(toptier_code) 

167 

168 # Make API request with params 

169 response = self._client._make_request( 

170 "GET", 

171 endpoint, 

172 params=params 

173 ) 

174 

175 return response