fix: mypy issue

Author: 1yam

tests/unit/services/test_pricing.py

@@ -10,6 +10,7 @@
     PAYG_GROUP,
     PRICING_GROUPS,
     GroupEntity,
+    Price,
     Pricing,
     PricingEntity,
     PricingModel,
@@ -64,17 +65,21 @@ async def test_get_pricing_aggregate(mock_client):
     storage_entity = result[PricingEntity.STORAGE]
     assert isinstance(storage_entity, PricingPerEntity)
     assert "storage" in storage_entity.price
-    assert storage_entity.price["storage"].holding == Decimal("0.333333333")
+    storage_price = storage_entity.price["storage"]
+    assert isinstance(storage_price, Price)  # Add type assertion for mypy
+    assert storage_price.holding == Decimal("0.333333333")
 
     # Check program entity has correct compute unit details
     program_entity = result[PricingEntity.PROGRAM]
     assert isinstance(program_entity, PricingPerEntity)
+    assert program_entity.compute_unit is not None  # Ensure compute_unit is not None
     assert program_entity.compute_unit.vcpus == 1
     assert program_entity.compute_unit.memory_mib == 2048
     assert program_entity.compute_unit.disk_mib == 2048
 
     # Check tiers in instance entity
     instance_entity = result[PricingEntity.INSTANCE]
+    assert instance_entity.tiers is not None  # Ensure tiers is not None
     assert len(instance_entity.tiers) == 6
     assert instance_entity.tiers[0].id == "tier-1"
     assert instance_entity.tiers[0].compute_units == 1
@@ -96,11 +101,14 @@ async def test_get_pricing_for_services(mock_client):
     assert PricingEntity.INSTANCE not in result
 
     # Verify specific pricing data
-    assert result[PricingEntity.STORAGE].price["storage"].holding == Decimal(
-        "0.333333333"
-    )
-    assert result[PricingEntity.PROGRAM].price["compute_unit"].payg == Decimal("0.011")
-    assert result[PricingEntity.PROGRAM].price["compute_unit"].holding == Decimal("200")
+    storage_price = result[PricingEntity.STORAGE].price["storage"]
+    assert isinstance(storage_price, Price)  # Ensure it's a Price object
+    assert storage_price.holding == Decimal("0.333333333")
+
+    compute_price = result[PricingEntity.PROGRAM].price["compute_unit"]
+    assert isinstance(compute_price, Price)  # Ensure it's a Price object
+    assert compute_price.payg == Decimal("0.011")
+    assert compute_price.holding == Decimal("200")
 
     # Test Case 2: Using pre-fetched pricing aggregate
     pricing_info = await pricing_service.get_pricing_aggregate()
@@ -131,10 +139,12 @@ async def test_get_pricing_for_services(mock_client):
     gpu_result = await pricing_service.get_pricing_for_services(gpu_services)
     assert len(gpu_result) == 2
     # Check GPU models are present
-    assert (
-        gpu_result[PricingEntity.INSTANCE_GPU_STANDARD].tiers[0].model == "RTX 4000 ADA"
-    )
-    assert gpu_result[PricingEntity.INSTANCE_GPU_PREMIUM].tiers[1].model == "H100"
+    standard_tiers = gpu_result[PricingEntity.INSTANCE_GPU_STANDARD].tiers
+    premium_tiers = gpu_result[PricingEntity.INSTANCE_GPU_PREMIUM].tiers
+    assert standard_tiers is not None
+    assert premium_tiers is not None
+    assert standard_tiers[0].model == "RTX 4000 ADA"
+    assert premium_tiers[1].model == "H100"
 
 
 @pytest.mark.asyncio
@@ -156,17 +166,27 @@ async def test_get_pricing_for_gpu_services(mock_client):
 
     # Verify GPU standard pricing and details
     gpu_standard = result[PricingEntity.INSTANCE_GPU_STANDARD]
-    assert gpu_standard.price["compute_unit"].payg == Decimal("0.28")
-    assert len(gpu_standard.tiers) == 5
-    assert gpu_standard.tiers[0].model == "RTX 4000 ADA"
-    assert gpu_standard.tiers[0].vram == 20480
+    compute_unit_price = gpu_standard.price["compute_unit"]
+    assert isinstance(compute_unit_price, Price)
+    assert compute_unit_price.payg == Decimal("0.28")
+
+    standard_tiers = gpu_standard.tiers
+    assert standard_tiers is not None
+    assert len(standard_tiers) == 5
+    assert standard_tiers[0].model == "RTX 4000 ADA"
+    assert standard_tiers[0].vram == 20480
 
     # Verify GPU premium pricing and details
     gpu_premium = result[PricingEntity.INSTANCE_GPU_PREMIUM]
-    assert gpu_premium.price["compute_unit"].payg == Decimal("0.56")
-    assert len(gpu_premium.tiers) == 2
-    assert gpu_premium.tiers[1].model == "H100"
-    assert gpu_premium.tiers[1].vram == 81920
+    premium_compute_price = gpu_premium.price["compute_unit"]
+    assert isinstance(premium_compute_price, Price)
+    assert premium_compute_price.payg == Decimal("0.56")
+
+    premium_tiers = gpu_premium.tiers
+    assert premium_tiers is not None
+    assert len(premium_tiers) == 2
+    assert premium_tiers[1].model == "H100"
+    assert premium_tiers[1].vram == 81920
 
 
 @pytest.mark.asyncio