From e3e91e6238d7f294095ca564f783ebc4ad6e412a Mon Sep 17 00:00:00 2001
From: Tom White <tom.e.white@gmail.com>
Date: Fri, 8 Mar 2024 13:32:04 +0000
Subject: [PATCH] Fix Modal deprecations (#417)

---
 cubed/runtime/executors/modal.py        |  7 ++++---
 cubed/tests/runtime/test_modal_async.py | 14 +++++++++++---
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/cubed/runtime/executors/modal.py b/cubed/runtime/executors/modal.py
index 5cbbf274..738764e0 100644
--- a/cubed/runtime/executors/modal.py
+++ b/cubed/runtime/executors/modal.py
@@ -70,7 +70,7 @@ def check_runtime_memory(spec):
 
 @stub.function(
     image=aws_image,
-    secret=modal.Secret.from_name("my-aws-secret"),
+    secrets=[modal.Secret.from_name("my-aws-secret")],
     memory=RUNTIME_MEMORY_MIB,
     retries=2,
     cloud="aws",
@@ -85,13 +85,14 @@ def run_remotely(input, func=None, config=None, name=None, compute_id=None):
 # For GCP we need to use a class so we can set up credentials by hooking into the container lifecycle
 @stub.cls(
     image=gcp_image,
-    secret=modal.Secret.from_name("my-googlecloud-secret"),
+    secrets=[modal.Secret.from_name("my-googlecloud-secret")],
     memory=RUNTIME_MEMORY_MIB,
     retries=2,
     cloud="gcp",
 )
 class Container:
-    def __enter__(self):
+    @modal.enter()
+    def set_up_credentials(self):
         json = os.environ["SERVICE_ACCOUNT_JSON"]
         path = os.path.abspath("application_credentials.json")
         with open(path, "w") as f:
diff --git a/cubed/tests/runtime/test_modal_async.py b/cubed/tests/runtime/test_modal_async.py
index 492669ae..3b4b13ec 100644
--- a/cubed/tests/runtime/test_modal_async.py
+++ b/cubed/tests/runtime/test_modal_async.py
@@ -34,19 +34,27 @@
 
 
 @stub.function(
-    image=image, secret=modal.Secret.from_name("my-aws-secret"), retries=2, timeout=10
+    image=image,
+    secrets=[modal.Secret.from_name("my-aws-secret")],
+    retries=2,
+    timeout=10,
 )
 def deterministic_failure_modal(i, path=None, timing_map=None, *, name=None):
     return deterministic_failure(path, timing_map, i, name=name)
 
 
-@stub.function(image=image, secret=modal.Secret.from_name("my-aws-secret"), timeout=10)
+@stub.function(
+    image=image, secrets=[modal.Secret.from_name("my-aws-secret")], timeout=10
+)
 def deterministic_failure_modal_no_retries(i, path=None, timing_map=None, *, name=None):
     return deterministic_failure(path, timing_map, i, name=name)
 
 
 @stub.function(
-    image=image, secret=modal.Secret.from_name("my-aws-secret"), retries=2, timeout=300
+    image=image,
+    secrets=[modal.Secret.from_name("my-aws-secret")],
+    retries=2,
+    timeout=300,
 )
 def deterministic_failure_modal_long_timeout(
     i, path=None, timing_map=None, *, name=None