fix: repair merge corruption in live_serverless and its tests

live_serverless.py was left syntactically/structurally broken by a bad
conflict resolution: an orphaned second docstring, a _live_image property
referencing an unimported DEFAULT_PYTHON_VERSION, the two template-injection
methods dedented into a dead nested scope, and unreachable code after each
validator's return.

Reconstructed to the validator-helper design (matches the latest commit's
intent): LiveServerlessMixin keeps _create_new_template /
_configure_existing_template as real methods that inject dockerArgs, and each
subclass validator defers to _apply_default_live_image(data, image_type). The
no-op _live_image property and _image_type ClassVars are removed.

Also fixes pre-existing breakage the reconstruction surfaced:
- test_live_serverless: assert the default image via get_image_name instead
  of the removed _live_image property
- test_live_load_balancer: import DEFAULT_PYTHON_VERSION (local_python_version
  is not defined in constants on this branch)
- test_cpu_disk_sizing: drop a body-less duplicate test (unused-var lint)
- test_lb_remote_execution: modernize the scanner fixture to @endpoint and
  rename it off the test_*.py prefix the scanner skips (was discovering 0)

Author: deanq

src/runpod_flash/core/resources/live_serverless.py

@@ -1,14 +1,9 @@
 # Ship serverless code as you write it. No builds, no deploys -- just run.
-from typing import Any, ClassVar
+from typing import Any
 
-# Ship serverless code as you write it. No builds, no deploys — just run.
 from pydantic import model_validator
 
-from .constants import (
-    GPU_BASE_IMAGE_PYTHON_VERSION,
-    get_image_name,
-    local_python_version,
-)
+from .constants import DEFAULT_PYTHON_VERSION, get_image_name
 from .injection import build_injection_cmd
 from .load_balancer_sls_resource import (
     CpuLoadBalancerSlsResource,
@@ -22,31 +17,30 @@
 class LiveServerlessMixin:
     """Common mixin for live serverless endpoints.
 
-    Treats the Flash runtime image as a *default*: if the caller passes an
-    ``imageName`` (e.g. via ``Endpoint(image=...)`` in client mode), that
-    value wins. Otherwise the Flash runtime image for this resource type is
-    used so decorator-mode workloads continue to deploy the Flash wrapper.
-
-    The default is applied via the ``@model_validator(mode="before")`` on each
-    concrete subclass (see ``_apply_default_live_image``); reads and writes of
-    ``imageName`` go through the normal Pydantic field machinery so model
-    serialization, drift detection, and ``setattr`` all stay consistent.
-    """Configures process injection via dockerArgs for any base image.
-
-    Sets a default base image (user can override via imageName) and generates
-    dockerArgs to download, extract, and run the flash-worker tarball at container
-    start time. QB vs LB mode is determined by FLASH_ENDPOINT_TYPE env var at
-    runtime, not by the Docker image.
+    Configures process injection via ``dockerArgs`` for any base image, and
+    treats the Flash runtime image as a *default*: if the caller passes an
+    ``imageName`` (e.g. via ``Endpoint(image=...)`` in client mode), that value
+    wins. Otherwise the Flash runtime image for this resource type is applied by
+    the ``@model_validator(mode="before")`` on each concrete subclass (see
+    ``_apply_default_live_image``), so decorator-mode workloads continue to
+    deploy the Flash wrapper.
+
+    The injection ``dockerArgs`` download, extract, and run the flash-worker
+    tarball at container start; QB vs LB mode is determined by the
+    ``FLASH_ENDPOINT_TYPE`` env var at runtime, not by the Docker image.
     """
 
-    _image_type: ClassVar[str] = (
-        ""  # override in subclasses: 'gpu', 'cpu', 'lb', 'lb-cpu'
-    )
+    def _create_new_template(self) -> PodTemplate:
+        """Create template with dockerArgs for process injection."""
+        template = super()._create_new_template()  # type: ignore[misc]
+        template.dockerArgs = build_injection_cmd()
+        return template
 
-    @property
-    def _live_image(self) -> str:
-        python_version = getattr(self, "python_version", None) or DEFAULT_PYTHON_VERSION
-        return get_image_name(self._image_type, python_version)
+    def _configure_existing_template(self) -> None:
+        """Configure existing template, adding dockerArgs for injection if not user-set."""
+        super()._configure_existing_template()  # type: ignore[misc]
+        if self.template is not None and not self.template.dockerArgs:  # type: ignore[attr-defined]
+            self.template.dockerArgs = build_injection_cmd()  # type: ignore[attr-defined]
 
 
 def _apply_default_live_image(data: Any, image_type: str):
@@ -63,18 +57,6 @@ def _apply_default_live_image(data: Any, image_type: str):
         data["imageName"] = get_image_name(image_type, python_version)
     return data
 
-    def _create_new_template(self) -> PodTemplate:
-        """Create template with dockerArgs for process injection."""
-        template = super()._create_new_template()  # type: ignore[misc]
-        template.dockerArgs = build_injection_cmd()
-        return template
-
-    def _configure_existing_template(self) -> None:
-        """Configure existing template, adding dockerArgs for injection if not user-set."""
-        super()._configure_existing_template()  # type: ignore[misc]
-        if self.template is not None and not self.template.dockerArgs:  # type: ignore[attr-defined]
-            self.template.dockerArgs = build_injection_cmd()  # type: ignore[attr-defined]
-
 
 class LiveServerless(LiveServerlessMixin, ServerlessEndpoint):
     """GPU-only live serverless endpoint."""
@@ -84,11 +66,6 @@ class LiveServerless(LiveServerlessMixin, ServerlessEndpoint):
     def set_live_serverless_template(cls, data: dict):
         """Default to the GPU Flash runtime image when none is supplied."""
         return _apply_default_live_image(data, "gpu")
-        """Set default GPU image for Live Serverless."""
-        if "imageName" not in data:
-            python_version = data.get("python_version") or GPU_BASE_IMAGE_PYTHON_VERSION
-            data["imageName"] = get_image_name("gpu", python_version)
-        return data
 
 
 class CpuLiveServerless(LiveServerlessMixin, CpuServerlessEndpoint):
@@ -99,11 +76,6 @@ class CpuLiveServerless(LiveServerlessMixin, CpuServerlessEndpoint):
     def set_live_serverless_template(cls, data: dict):
         """Default to the CPU Flash runtime image when none is supplied."""
         return _apply_default_live_image(data, "cpu")
-        """Set default CPU image for Live Serverless."""
-        if "imageName" not in data:
-            python_version = data.get("python_version") or local_python_version()
-            data["imageName"] = get_image_name("cpu", python_version)
-        return data
 
 
 class LiveLoadBalancer(LiveServerlessMixin, LoadBalancerSlsResource):
@@ -114,11 +86,6 @@ class LiveLoadBalancer(LiveServerlessMixin, LoadBalancerSlsResource):
     def set_live_lb_template(cls, data: dict):
         """Default to the LB Flash runtime image when none is supplied."""
         return _apply_default_live_image(data, "lb")
-        """Set default image for Live Load-Balanced endpoint."""
-        if "imageName" not in data:
-            python_version = data.get("python_version") or GPU_BASE_IMAGE_PYTHON_VERSION
-            data["imageName"] = get_image_name("lb", python_version)
-        return data
 
 
 class CpuLiveLoadBalancer(LiveServerlessMixin, CpuLoadBalancerSlsResource):
@@ -129,8 +96,3 @@ class CpuLiveLoadBalancer(LiveServerlessMixin, CpuLoadBalancerSlsResource):
     def set_live_cpu_lb_template(cls, data: dict):
         """Default to the CPU LB Flash runtime image when none is supplied."""
         return _apply_default_live_image(data, "lb-cpu")
-        """Set default CPU image for Live Load-Balanced endpoint."""
-        if "imageName" not in data:
-            python_version = data.get("python_version") or local_python_version()
-            data["imageName"] = get_image_name("lb-cpu", python_version)
-        return data

tests/integration/test_cpu_disk_sizing.py

@@ -244,15 +244,6 @@ def test_mixed_cpu_generations_integration(self):
         assert "cpu5c-1-2: max 15GB" in error_msg
 
 
-class TestLiveServerlessImageIntegration:
-    """Test image default + override behavior in live serverless variants (AE-3153)."""
-
-    def test_live_serverless_image_consistency(self):
-        """LiveServerless variants default to distinct Flash runtime images."""
-        gpu_live = LiveServerless(name="gpu-live")
-        cpu_live = CpuLiveServerless(name="cpu-live")
-
-        # Verify different default images are used per resource type.
 class TestLiveServerlessImageDefaultsIntegration:
     """Test image defaults in live serverless variants."""
 

tests/integration/test_lb_remote_execution.py

@@ -132,6 +132,7 @@ def test_live_load_balancer_image_default_and_override(self):
         # Guard against a future regression where both paths collapse to the
         # same default (e.g. the override branch reverting to a no-op).
         assert default_lb.imageName != custom_lb.imageName
+
     def test_live_load_balancer_default_image(self):
         """Test that LiveLoadBalancer uses GPU LB base image by default."""
         lb = LiveLoadBalancer(name="test-api")
@@ -170,50 +171,50 @@ async def qb_func():
         assert qb_func.__remote_config__["path"] is None
 
     def test_scanner_discovers_load_balancer_resources(self):
-        """Test that scanner can discover LiveLoadBalancer and LoadBalancerSlsResource."""
+        """Test that scanner discovers @Endpoint load-balanced route handlers."""
         from runpod_flash.cli.commands.build_utils.scanner import RuntimeScanner
         from pathlib import Path
         import tempfile
 
-        # Create temporary Python file with LoadBalancer resource
+        # Uses the current @Endpoint API (remote is deprecated). The worker file
+        # must NOT be named test_*.py — the scanner skips test files, so a
+        # test-prefixed fixture would yield zero discovered functions.
         code = """
-from runpod_flash import LiveLoadBalancer, LoadBalancerSlsResource, remote
+from runpod_flash.endpoint import Endpoint
+from runpod_flash.core.resources.gpu import GpuGroup
 
-# Test LiveLoadBalancer discovery
-api = LiveLoadBalancer(name="test-api")
+api = Endpoint(name="test-api", gpu=GpuGroup.AMPERE_16)
 
-@remote(api, method="POST", path="/api/process")
+@api.post("/api/process")
 async def process_data(x: int):
     return {"result": x}
 
-# Test LoadBalancerSlsResource discovery
-deployed = LoadBalancerSlsResource(name="deployed-api", imageName="test:latest")
+status_api = Endpoint(name="deployed-api", gpu=GpuGroup.AMPERE_16)
 
-@remote(deployed, method="GET", path="/api/status")
+@status_api.get("/api/status")
 def get_status():
     return {"status": "ok"}
 """
 
         with tempfile.TemporaryDirectory() as tmpdir:
             project_dir = Path(tmpdir)
-            py_file = project_dir / "test_api.py"
+            py_file = project_dir / "worker_api.py"
             py_file.write_text(code)
 
             scanner = RuntimeScanner(project_dir)
             functions = scanner.discover_remote_functions()
 
-            # Verify both resources were discovered
+            # Both LB route handlers are discovered.
             assert len(functions) == 2
 
-            # Verify resource types are correctly identified
-            resource_types = {f.resource_type for f in functions}
-            assert "LiveLoadBalancer" in resource_types
-            assert "LoadBalancerSlsResource" in resource_types
-
-            # Verify resource configs were extracted
-            assert "test-api-fb" in scanner.resource_types
-            assert scanner.resource_types["test-api-fb"] == "LiveLoadBalancer"
-            assert "deployed-api-fb" in scanner.resource_types
-            assert (
-                scanner.resource_types["deployed-api-fb"] == "LoadBalancerSlsResource"
-            )
+            # Endpoint(gpu=...) with route handlers resolves to LiveLoadBalancer.
+            assert {f.resource_type for f in functions} == {"LiveLoadBalancer"}
+            assert all(f.is_load_balanced for f in functions)
+
+            # HTTP method and path are captured per route.
+            routes = {(f.http_method, f.http_path) for f in functions}
+            assert routes == {("POST", "/api/process"), ("GET", "/api/status")}
+
+            # Resource configs are tracked by name.
+            assert scanner.resource_types["test-api"] == "LiveLoadBalancer"
+            assert scanner.resource_types["deployed-api"] == "LiveLoadBalancer"

tests/unit/resources/test_live_load_balancer.py

@@ -5,8 +5,8 @@
 
 import pytest
 from runpod_flash.core.resources.constants import (
+    DEFAULT_PYTHON_VERSION,
     GPU_BASE_IMAGE_PYTHON_VERSION,
-    local_python_version,
 )
 from runpod_flash.core.resources.cpu import CpuInstanceType
 from runpod_flash.core.resources.live_serverless import (
@@ -202,7 +202,7 @@ def test_cpu_live_load_balancer_default_image_tag(self):
         os.environ.pop("FLASH_IMAGE_TAG", None)
 
         lb = CpuLiveLoadBalancer(name="test-lb")
-        assert f"py{local_python_version()}" in lb.imageName
+        assert f"py{DEFAULT_PYTHON_VERSION}" in lb.imageName
         assert lb.template is not None
         assert lb.template.imageName == lb.imageName
 

tests/unit/resources/test_live_serverless.py

@@ -53,6 +53,7 @@ def test_live_serverless_image_default_unchanged(self):
         """LiveServerless still defaults to the Flash GPU runtime image."""
         live_serverless = LiveServerless(name="example_gpu_live_serverless")
         assert "flash:" in live_serverless.imageName
+
     def test_live_serverless_user_can_override_image(self):
         """Test user can set custom imageName (BYOI)."""
         live_serverless = LiveServerless(
@@ -128,6 +129,7 @@ def test_cpu_live_serverless_image_default_unchanged(self):
             instanceIds=[CpuInstanceType.CPU3G_1_4],
         )
         assert "flash-cpu:" in live_serverless.imageName
+
     def test_cpu_live_serverless_user_can_override_image(self):
         """Test CpuLiveServerless allows user to set custom image."""
         live_serverless = CpuLiveServerless(name="test", imageName="python:3.11-slim")
@@ -200,13 +202,28 @@ def test_docker_args_set_on_existing_template(self):
 
     def test_image_name_property_gpu(self):
         """LiveServerless defaults imageName to the Flash runtime image when none supplied."""
+        from runpod_flash.core.resources.constants import (
+            DEFAULT_PYTHON_VERSION,
+            get_image_name,
+        )
+
         live_serverless = LiveServerless(name="test")
-        assert live_serverless.imageName == live_serverless._live_image
+        assert live_serverless.imageName == get_image_name(
+            "gpu", DEFAULT_PYTHON_VERSION
+        )
 
     def test_image_name_property_cpu(self):
         """CpuLiveServerless defaults imageName to the Flash runtime image when none supplied."""
+        from runpod_flash.core.resources.constants import (
+            DEFAULT_PYTHON_VERSION,
+            get_image_name,
+        )
+
         live_serverless = CpuLiveServerless(name="test")
-        assert live_serverless.imageName == live_serverless._live_image
+        assert live_serverless.imageName == get_image_name(
+            "cpu", DEFAULT_PYTHON_VERSION
+        )
+
     def test_all_live_classes_have_docker_args(self):
         """Test all Live* classes set dockerArgs on their templates."""
         classes_and_kwargs = [
@@ -262,6 +279,7 @@ def test_default_image_validator_passes_through_non_dict(self):
         original = LiveServerless(name="test", imageName="byo/image:v1")
         revalidated = LiveServerless.model_validate(original)
         assert revalidated.imageName == "byo/image:v1"
+
     def test_live_serverless_byoi_gpu(self):
         """Test LiveServerless respects user-provided imageName."""
         live_serverless = LiveServerless(name="test", imageName="custom/gpu:v1")