From 346f8bf4f6b3256b3ffd4854af05990ce40f53fa Mon Sep 17 00:00:00 2001
From: Bernhard Mallinger <bernhard.mallinger@eox.at>
Date: Wed, 30 Oct 2024 08:54:53 +0100
Subject: [PATCH] Allow customizing resource requirements for s3fs

---
 pygeoapi_kubernetes_papermill/common.py | 11 ++++++++---
 tests/test_notebook_processor.py        | 18 +++++++++++++++++-
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/pygeoapi_kubernetes_papermill/common.py b/pygeoapi_kubernetes_papermill/common.py
index 81972a5..4fc1d00 100644
--- a/pygeoapi_kubernetes_papermill/common.py
+++ b/pygeoapi_kubernetes_papermill/common.py
@@ -167,6 +167,8 @@ def extra_configs() -> Iterable[ExtraConfig]:
                     secret_name=self.s3["secret_name"],
                     mount_path=self.s3["mount_path"],
                     s3_url=self.s3["s3_url"],
+                    resource_limits=self.s3["resource_limits"],
+                    resource_requests=self.s3["resource_requests"],
                 )
 
             access_functions = {
@@ -210,7 +212,9 @@ def build(input_dict: dict):
     )
 
 
-def s3_config(bucket_name, secret_name, s3_url, mount_path) -> ExtraConfig:
+def s3_config(
+    bucket_name, secret_name, s3_url, mount_path, resource_requests, resource_limits
+) -> ExtraConfig:
     s3_user_bucket_volume_name = "s3-user-bucket"
     return ExtraConfig(
         volume_mounts=[
@@ -265,11 +269,12 @@ def s3_config(bucket_name, secret_name, s3_url, mount_path) -> ExtraConfig:
                     ),
                 ],
                 resources=k8s_client.V1ResourceRequirements(
-                    limits={"cpu": "0.2", "memory": "512Mi"},
+                    limits={"cpu": "0.2", "memory": "512Mi"} | resource_limits,
                     requests={
                         "cpu": "0.05",
                         "memory": "32Mi",
-                    },
+                    }
+                    | resource_requests,
                 ),
                 env=[
                     k8s_client.V1EnvVar(name="S3FS_ARGS", value="-oallow_other"),
diff --git a/tests/test_notebook_processor.py b/tests/test_notebook_processor.py
index 7127358..d7c38b6 100644
--- a/tests/test_notebook_processor.py
+++ b/tests/test_notebook_processor.py
@@ -190,7 +190,15 @@ def test_no_s3_bucket_by_default(papermill_processor, create_pod_kwargs):
 @pytest.fixture()
 def papermill_processor_s3():
     return _create_processor(
-        {"s3": {"bucket_name": "example", "secret_name": "example", "s3_url": ""}}
+        {
+            "s3": {
+                "bucket_name": "example",
+                "secret_name": "example",
+                "s3_url": "",
+                "resource_requests": {"ice/cream": 1},
+                "resource_limits": {"ice/cream": 3},
+            }
+        }
     )
 
 
@@ -641,6 +649,14 @@ def test_extra_requirements_are_added(create_pod_kwargs):
     assert resources.limits["ice/cream"] == 3
 
 
+def test_s3fs_requirements_are_added(create_pod_kwargs, papermill_processor_s3):
+    job_pod_spec = papermill_processor_s3.create_job_pod_spec(**create_pod_kwargs)
+
+    resources = job_pod_spec.pod_spec.containers[1].resources
+    assert resources.requests["ice/cream"] == 1
+    assert resources.limits["ice/cream"] == 3
+
+
 def test_invalid_params_raises_user_error(papermill_processor, create_pod_kwargs_with):
     with pytest.raises(ProcessorClientError, match=".*mem_limit.*"):
         papermill_processor.create_job_pod_spec(