Coverage for src/usaspending/models/award_factory.py: 100%

1"""Factory for creating appropriate Award subclass instances.""" 

2 

3from __future__ import annotations 

4from typing import Dict, Any, Optional, TYPE_CHECKING 

5 

6from ..config import CONTRACT_CODES, IDV_CODES, GRANT_CODES, LOAN_CODES 

7from ..exceptions import ValidationError 

8 

9if TYPE_CHECKING: 

10 from ..client import USASpending 

11 from .award import Award 

12 

13 

14def create_award( 

15 data_or_id: Dict[str, Any] | str, client: Optional[USASpending] = None 

16) -> Award: 

17 """ 

18 Factory function to create the appropriate Award subclass based on the award data. 

19 

20 Args: 

21 data_or_id: Award data dictionary or unique award ID string 

22 client: Optional USASpending client instance 

23 

24 Returns: 

25 Appropriate Award subclass instance (Contract, Grant, IDV, Loan, or base Award) 

26 """ 

27 # Import here to avoid circular imports 

28 from .award import Award 

29 from .contract import Contract 

30 from .grant import Grant 

31 from .idv import IDV 

32 from .loan import Loan 

33 

34 # If it's just an ID, create base Award and let lazy loading determine type 

35 if isinstance(data_or_id, str): 

36 return Award(data_or_id, client) 

37 

38 if not isinstance(data_or_id, dict): 

39 raise ValidationError("Award factory expects a dict or an award_id string") 

40 

41 # Determine award type from data 

42 category = data_or_id.get("category", "").lower() 

43 award_type = data_or_id.get("type", "") 

44 

45 award_class_map = { 

46 "contract": Contract, 

47 "idv": IDV, 

48 "grant": Grant, 

49 "loan": Loan, 

50 } 

51 

52 # Determine class from category first (most reliable) 

53 award_class = award_class_map.get(category) 

54 

55 # Fallback to type codes if category doesn't match 

56 if not award_class: 

57 if award_type in CONTRACT_CODES: 

58 award_class = Contract 

59 elif award_type in IDV_CODES: 

60 award_class = IDV 

61 elif award_type in GRANT_CODES: 

62 award_class = Grant 

63 elif award_type in LOAN_CODES: 

64 award_class = Loan 

65 

66 # Use the determined class or default to the base Award class 

67 cls = award_class or Award 

68 return cls(data_or_id, client)