first commit

tests/test_inventory.py

@@ -0,0 +1,57 @@
+from gpu_rent.inventory import (
+    gpu_boot_image_score,
+    gpu_quota_from_compute,
+    match_label,
+    pick_boot_image,
+    rank_flavors,
+)
+
+
+class FakeFlavor:
+    def __init__(self, flavor_id: str, name: str, extra=None, disabled: bool = False):
+        self.id = flavor_id
+        self.name = name
+        self.extra_specs = extra or {}
+        self.is_disabled = disabled
+        self.vcpus = 8
+        self.ram = 32768
+
+
+def test_match_4090_24_not_48():
+    a = FakeFlavor("1", "GPU 1x RTX 4090 24GB")
+    b = FakeFlavor("2", "GPU 1x RTX 4090 48GB")
+    assert match_label("4090-24", a)
+    assert not match_label("4090-24", b)
+    assert match_label("4090-48", b)
+
+
+def test_rank_skips_disabled():
+    flavors = [
+        FakeFlavor("d", "RTX 4090 24GB", disabled=True),
+        FakeFlavor("ok", "RTX A5000 24GB"),
+    ]
+    ranked = rank_flavors(flavors, ("4090-24", "a5000"))
+    assert [x.id for x in ranked] == ["ok"]
+    assert ranked[0].label == "a5000"
+
+
+def test_gpu_quota_from_compute():
+    assert gpu_quota_from_compute({"cores": 10}) is None
+    assert gpu_quota_from_compute({"gpu": 0}) == 0
+    assert gpu_quota_from_compute({"GPU_limit": 2}) == 2
+
+
+def test_pick_boot_image_prefers_24_580_without_docker():
+    class Img:
+        def __init__(self, name):
+            self.name = name
+
+    images = [
+        Img("Ubuntu 24.04 LTS 64-bit GPU Driver 580 Docker"),
+        Img("Ubuntu 24.04 LTS 64-bit GPU Driver 535"),
+        Img("Ubuntu 24.04 LTS 64-bit GPU Driver 580"),
+        Img("Data Science VM (Ubuntu 22.04 LTS 64-bit)"),
+    ]
+    picked = pick_boot_image(images)
+    assert picked.name == "Ubuntu 24.04 LTS 64-bit GPU Driver 580"
+    assert gpu_boot_image_score(picked.name) > gpu_boot_image_score(images[0].name)