From cc6244154e4feca918b4afe4e8692b95d6e2a08f Mon Sep 17 00:00:00 2001
From: Chi Chang <chi@yelp.com>
Date: Wed, 21 Feb 2024 09:56:09 -0800
Subject: [PATCH] Add unit tests

---
 service_configuration_lib/utils.py | 13 +++++--
 tests/utils_test.py                | 55 ++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/service_configuration_lib/utils.py b/service_configuration_lib/utils.py
index 52725e2..fcfb6eb 100644
--- a/service_configuration_lib/utils.py
+++ b/service_configuration_lib/utils.py
@@ -17,11 +17,14 @@
 import yaml
 
 DEFAULT_SPARK_RUN_CONFIG = '/nail/srv/configs/spark.yaml'
-
 POD_TEMPLATE_PATH = '/nail/tmp/spark-pt-{file_uuid}.yaml'
-
 SPARK_EXECUTOR_POD_TEMPLATE = '/nail/srv/configs/spark_executor_pod_template.yaml'
 
+LOCALHOST = '127.0.0.1'
+EPHEMERAL_PORT_START = 49152
+EPHEMERAL_PORT_END = 65535
+
+
 log = logging.Logger(__name__)
 log.setLevel(logging.INFO)
 
@@ -46,7 +49,11 @@ def load_spark_srv_conf(preset_values=None) -> Tuple[Mapping, Mapping, Mapping,
         raise e
 
 
-def ephemeral_port_reserve_range(preferred_port_start: int, preferred_port_end: int, ip='127.0.0.1') -> int:
+def ephemeral_port_reserve_range(
+    preferred_port_start: int = EPHEMERAL_PORT_START,
+    preferred_port_end: int = EPHEMERAL_PORT_END,
+    ip: str = LOCALHOST,
+) -> int:
     """
     Pick an available from the preferred port range. If all ports from the port range are unavailable,
     pick a random available ephemeral port.
diff --git a/tests/utils_test.py b/tests/utils_test.py
index ff2c8b6..6db80a3 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -1,9 +1,52 @@
 import uuid
+from socket import SO_REUSEADDR
+from socket import socket as Socket
+from socket import SOL_SOCKET
 from unittest import mock
+from unittest.mock import mock_open
+from unittest.mock import patch
 
 import pytest
 
 from service_configuration_lib import utils
+from service_configuration_lib.utils import ephemeral_port_reserve_range
+from service_configuration_lib.utils import LOCALHOST
+
+
+MOCK_ENV_NAME = 'mock_env_name'
+
+
+def _bind_reuse(ip, port):
+    sock = Socket()
+    sock.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
+    sock.bind((ip, port))
+    return sock
+
+
+def test_preferred_port():
+    port = ephemeral_port_reserve_range()
+    port2 = ephemeral_port_reserve_range(port, port + 1)
+    assert port == port2
+    assert _bind_reuse(LOCALHOST, port2)
+
+
+def test_preferred_port_in_use():
+    """if preferred port is in use, it will find an unused port"""
+    port = ephemeral_port_reserve_range()
+    sock = _bind_reuse(LOCALHOST, port)
+    sock.listen(1)  # make the port in-use
+
+    port_end = port + 10
+    port2 = ephemeral_port_reserve_range(port, port_end)
+    assert port != port2
+    assert port2 > port and port2 <= port_end
+    assert _bind_reuse(LOCALHOST, port2)
+
+
+def test_get_random_string():
+    length = 50
+    result = utils.get_random_string(50)
+    assert len(result) == length
 
 
 @pytest.mark.parametrize(
@@ -29,3 +72,15 @@ def test_get_k8s_resource_name_limit_size_with_hash(instance_name, expected_inst
 def test_generate_pod_template_path(hex_value):
     with mock.patch.object(uuid, 'uuid4', return_value=uuid.UUID(hex=hex_value)):
         assert utils.generate_pod_template_path() == f'/nail/tmp/spark-pt-{hex_value}.yaml'
+
+
+@pytest.fixture
+def mock_runtimeenv():
+    with patch('builtins.open', mock_open(read_data=MOCK_ENV_NAME)) as m:
+        yield m
+
+
+def test_get_runtime_env(mock_runtimeenv):
+    result = utils.get_runtime_env()
+    assert result == MOCK_ENV_NAME
+    mock_runtimeenv.assert_called_once_with('/nail/etc/runtimeenv', mode='r')