From 6b2d2a21fec46218b64cb62e64fad67f82d09f92 Mon Sep 17 00:00:00 2001
From: Miguel Angel Ajo Pelayo <majopela@redhat.com>
Date: Wed, 8 Jan 2025 10:59:49 +0100
Subject: [PATCH] Fix JMP_LEASE handling in pytest

1) Look for JMP_LEASE inside the client.lease methods
2) Only release runtime created leases, leases received via JMP_LEASE
   are kept, because those must be manually released.
---
 jumpstarter/cli/client.py     | 12 ++----------
 jumpstarter/config/client.py  | 18 +++++++++++-------
 jumpstarter/testing/pytest.py |  6 +-----
 3 files changed, 14 insertions(+), 22 deletions(-)

diff --git a/jumpstarter/cli/client.py b/jumpstarter/cli/client.py
index 5f66b23..f735fc0 100644
--- a/jumpstarter/cli/client.py
+++ b/jumpstarter/cli/client.py
@@ -1,12 +1,11 @@
 import logging
-import os
 from typing import Optional
 
 import click
 
 from jumpstarter.common import MetadataFilter
 from jumpstarter.common.utils import launch_shell
-from jumpstarter.config import ClientConfigV1Alpha1, ClientConfigV1Alpha1Drivers, UserConfigV1Alpha1, env
+from jumpstarter.config import ClientConfigV1Alpha1, ClientConfigV1Alpha1Drivers, UserConfigV1Alpha1
 
 from .util import AliasedGroup, make_table
 from .version import version
@@ -223,14 +222,7 @@ def client_shell(name: str, labels, lease_name):
     if not config:
         raise ValueError("no client specified")
 
-    # lease_name can be provided as an argument or via environment variable
-    lease_name = lease_name or os.environ.get(env.JMP_LEASE, "")
-
-    # when no lease name is provided, release the lease on exit
-    release_lease = lease_name == ""
-
-    with config.lease(metadata_filter=MetadataFilter(labels=dict(labels)), lease_name=lease_name,
-                      release=release_lease) as lease:
+    with config.lease(metadata_filter=MetadataFilter(labels=dict(labels)), lease_name=lease_name) as lease:
         with lease.serve_unix() as path:
             launch_shell(path, config.drivers.allow, config.drivers.unsafe)
 
diff --git a/jumpstarter/config/client.py b/jumpstarter/config/client.py
index 2207e50..1b8090a 100644
--- a/jumpstarter/config/client.py
+++ b/jumpstarter/config/client.py
@@ -13,7 +13,7 @@
 from jumpstarter.v1 import jumpstarter_pb2, jumpstarter_pb2_grpc
 
 from .common import CONFIG_PATH
-from .env import JMP_DRIVERS_ALLOW, JMP_ENDPOINT, JMP_TOKEN
+from .env import JMP_DRIVERS_ALLOW, JMP_ENDPOINT, JMP_LEASE, JMP_TOKEN
 from .tls import TLSConfigV1Alpha1
 
 
@@ -56,10 +56,10 @@ async def channel(self):
         return aio_secure_channel(self.endpoint, credentials)
 
     @contextmanager
-    def lease(self, metadata_filter: MetadataFilter, lease_name: str | None, release: bool = False):
+    def lease(self, metadata_filter: MetadataFilter, lease_name: str | None):
         with start_blocking_portal() as portal:
             with portal.wrap_async_context_manager(
-                self.lease_async(metadata_filter, lease_name, portal, release)) as lease:
+                self.lease_async(metadata_filter, lease_name, portal)) as lease:
                 yield lease
 
     def request_lease(self, metadata_filter: MetadataFilter):
@@ -97,10 +97,14 @@ async def release_lease_async(self, name):
         await controller.ReleaseLease(jumpstarter_pb2.ReleaseLeaseRequest(name=name))
 
     @asynccontextmanager
-    async def lease_async(self, metadata_filter: MetadataFilter, lease_name: str | None, portal: BlockingPortal,
-                          release=True):
-        # dynamically import to avoid circular imports
+    async def lease_async(self, metadata_filter: MetadataFilter, lease_name: str | None, portal: BlockingPortal):
         from jumpstarter.client import Lease
+
+        # if no lease_name provided, check if it is set in the environment
+        lease_name = lease_name or os.environ.get(JMP_LEASE, "")
+        # when no lease name is provided, release the lease on exit
+        release_lease = lease_name == ""
+
         async with Lease(
             channel=await self.channel(),
             name=lease_name,
@@ -108,7 +112,7 @@ async def lease_async(self, metadata_filter: MetadataFilter, lease_name: str | N
             portal=portal,
             allow=self.drivers.allow,
             unsafe=self.drivers.unsafe,
-            release=release,
+            release=release_lease,
             tls_config=self.tls,
         ) as lease:
             yield lease
diff --git a/jumpstarter/testing/pytest.py b/jumpstarter/testing/pytest.py
index 67c98b0..ebe71f9 100644
--- a/jumpstarter/testing/pytest.py
+++ b/jumpstarter/testing/pytest.py
@@ -1,5 +1,4 @@
 import logging
-import os
 import time
 from typing import ClassVar
 
@@ -59,10 +58,7 @@ def client(self):
         except RuntimeError:
             labels = getattr(self, "filter_labels", {})
             config = ClientConfigV1Alpha1.load("default")
-            lease_name = os.getenv(env.JMP_LEASE, None)
-            if lease_name is not None:
-                log.info(f"Using lease from env var JMP_LEASE: {lease_name}")
-            with config.lease(metadata_filter=MetadataFilter(labels=labels), lease_name=lease_name) as lease:
+            with config.lease(metadata_filter=MetadataFilter(labels=labels)) as lease:
                 with lease.connect() as client:
                     yield client
         # BUG workaround: make sure that grpc servers get the client/lease release properly