From 342961611a4c47b5379e95ec377ca6289c8d75be Mon Sep 17 00:00:00 2001
From: Thomas Carmet <8408330+tcarmet@users.noreply.github.com>
Date: Tue, 30 Jan 2024 08:19:19 -0800
Subject: [PATCH] PTFE-1402 enable or disable external IP GCP runner (#523)

---
 runner_manager/backend/gcloud.py | 20 ++++++++++----------
 runner_manager/models/backend.py |  1 +
 tests/unit/backend/test_gcp.py   |  8 +++++++-
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/runner_manager/backend/gcloud.py b/runner_manager/backend/gcloud.py
index 43cabab4..f3e19beb 100644
--- a/runner_manager/backend/gcloud.py
+++ b/runner_manager/backend/gcloud.py
@@ -91,16 +91,16 @@ def disks(self) -> List[AttachedDisk]:
 
     @property
     def network_interfaces(self) -> List[NetworkInterface]:
-        return [
-            NetworkInterface(
-                network=self.instance_config.network,
-                access_configs=[
-                    AccessConfig(
-                        name="External NAT",
-                    )
-                ],
-            )
-        ]
+        network_interface: NetworkInterface = NetworkInterface(
+            network=self.instance_config.network,
+        )
+        if self.instance_config.enable_external_ip:
+            network_interface.access_configs = [
+                AccessConfig(
+                    name="External NAT",
+                )
+            ]
+        return [network_interface]
 
     @property
     def scheduling(self) -> Scheduling:
diff --git a/runner_manager/models/backend.py b/runner_manager/models/backend.py
index 6224ea88..04d9658f 100644
--- a/runner_manager/models/backend.py
+++ b/runner_manager/models/backend.py
@@ -104,6 +104,7 @@ class GCPInstanceConfig(InstanceConfig):
     machine_type: str = "e2-small"
     network: str = "global/networks/default"
     enable_nested_virtualization: bool = True
+    enable_external_ip: bool = True
     spot: bool = False
     disk_size_gb: int = 20
     disk_type: Literal["pd-ssd", "pd-standard"] = "pd-ssd"
diff --git a/tests/unit/backend/test_gcp.py b/tests/unit/backend/test_gcp.py
index 4a6df0b4..c7b73e42 100644
--- a/tests/unit/backend/test_gcp.py
+++ b/tests/unit/backend/test_gcp.py
@@ -1,7 +1,7 @@
 import os
 from typing import List
 
-from google.cloud.compute import Image
+from google.cloud.compute import Image, NetworkInterface
 from pytest import fixture, mark, raises
 from redis_om import NotFoundError
 
@@ -57,6 +57,12 @@ def gcp_runner(runner: Runner, gcp_group: RunnerGroup) -> Runner:
 def test_gcp_network_interfaces(gcp_group: RunnerGroup):
     interfaces = gcp_group.backend.network_interfaces
     assert len(interfaces) == 1
+    assert interfaces[0].access_configs[0].name == "External NAT"
+    # Test disabling external IP
+    gcp_group.backend.instance_config.enable_external_ip = False
+    interfaces: List[NetworkInterface] = gcp_group.backend.network_interfaces
+    assert len(interfaces) == 1
+    assert len(interfaces[0].access_configs) == 0
 
 
 def test_gcp_group(gcp_group: RunnerGroup):