Coverage for polar/models/subscription_product_price.py: 55%

1from typing import TYPE_CHECKING, Self 

2from uuid import UUID 

3 

4from sqlalchemy import ForeignKey, Integer, Uuid 

5from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship 

6 

7from polar.kit.db.models import RecordModel 

8from polar.models.product_price import ( 

9 LegacyRecurringProductPriceCustom, 

10 LegacyRecurringProductPriceFixed, 

11 ProductPrice, 

12 ProductPriceCustom, 

13 ProductPriceFixed, 

14 ProductPriceSeatUnit, 

15) 

16 

17if TYPE_CHECKING: 17 ↛ 18line 17 didn't jump to line 18 because the condition on line 17 was never true

18 from polar.models import Subscription 

19 

20 

21class SubscriptionProductPrice(RecordModel): 

22 __tablename__ = "subscription_product_prices" 

23 

24 subscription_id: Mapped[UUID] = mapped_column( 

25 Uuid, 

26 ForeignKey("subscriptions.id", ondelete="cascade"), 

27 primary_key=True, 

28 ) 

29 product_price_id: Mapped[UUID] = mapped_column( 

30 Uuid, 

31 ForeignKey("product_prices.id", ondelete="restrict"), 

32 primary_key=True, 

33 ) 

34 amount: Mapped[int] = mapped_column(Integer, nullable=False) 

35 

36 @declared_attr 

37 def product_price(cls) -> Mapped["ProductPrice"]: 

38 # This is an association table, so eager loading makes sense 

39 return relationship("ProductPrice", lazy="joined") 

40 

41 @declared_attr 

42 def subscription(cls) -> Mapped["Subscription"]: 

43 return relationship("Subscription", lazy="raise_on_sql") 

44 

45 @classmethod 

46 def from_price( 

47 cls, 

48 price: "ProductPrice", 

49 amount: int | None = None, 

50 seats: int | None = None, 

51 ) -> Self: 

52 if isinstance(price, ProductPriceFixed | LegacyRecurringProductPriceFixed): 

53 amount = price.price_amount 

54 elif isinstance(price, ProductPriceCustom | LegacyRecurringProductPriceCustom): 

55 assert amount is not None, "amount must be provided for custom prices" 

56 elif isinstance(price, ProductPriceSeatUnit): 

57 assert seats is not None, "seats must be provided for seat-based prices" 

58 amount = price.calculate_amount(seats) 

59 else: 

60 amount = 0 

61 return cls(product_price=price, amount=amount)