172 lines
7.0 KiB
Python
172 lines
7.0 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
sys.path.insert(0, "/www/wwwroot/getDomain/domainCheck")
|
|
|
|
from app.config import config # noqa: E402
|
|
from app.utils.database import _local_config_search_roots, _resolve_db_pool_limits # noqa: E402
|
|
|
|
|
|
class DatabasePoolLimitTests(unittest.TestCase):
|
|
def _write_json(self, root: str, name: str, payload: dict) -> None:
|
|
Path(root, name).write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
|
|
|
|
def test_explicit_total_budget_override_keeps_previous_pool_shape(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir, patch.object(
|
|
config, "NODE_CODE", "mainland-controller-01-a"
|
|
), patch.dict(
|
|
os.environ,
|
|
{
|
|
"WORKER_PARENT_NODE_CODE": "mainland-controller-01",
|
|
"DOMAINCHECK_DB_POOL_TOTAL_BUDGET": "480",
|
|
"DB_POOL_SIZE": "",
|
|
"DB_POOL_WARM_SIZE": "",
|
|
"DB_POOL_IDLE_KEEP_MAX": "",
|
|
},
|
|
clear=False,
|
|
):
|
|
self._write_json(temp_dir, "thread_count.json", {"thread_count": "1000"})
|
|
self._write_json(temp_dir, "process_count.json", {"process_count": "80"})
|
|
self._write_json(temp_dir, "node_thread_counts.json", {"mainland-controller-01": 1200})
|
|
self._write_json(temp_dir, "node_process_counts.json", {"mainland-controller-01": 80})
|
|
|
|
previous_cwd = os.getcwd()
|
|
try:
|
|
os.chdir(temp_dir)
|
|
limits = _resolve_db_pool_limits()
|
|
finally:
|
|
os.chdir(previous_cwd)
|
|
|
|
self.assertEqual(6, limits["pool_size"])
|
|
self.assertEqual(1, limits["pool_warm_size"])
|
|
self.assertEqual(3, limits["pool_idle_keep_max"])
|
|
self.assertEqual(1200, limits["scaling_hints"]["thread_count"])
|
|
self.assertEqual(80, limits["scaling_hints"]["process_count"])
|
|
|
|
def test_pool_budget_scales_down_for_mid_sized_multi_process_workers(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir, patch.object(
|
|
config, "NODE_CODE", "mainland-worker-01-a"
|
|
), patch.dict(
|
|
os.environ,
|
|
{
|
|
"WORKER_PARENT_NODE_CODE": "mainland-worker-01",
|
|
"DB_POOL_SIZE": "",
|
|
"DB_POOL_WARM_SIZE": "",
|
|
"DB_POOL_IDLE_KEEP_MAX": "",
|
|
"DOMAINCHECK_DB_POOL_TOTAL_BUDGET": "",
|
|
},
|
|
clear=False,
|
|
):
|
|
self._write_json(temp_dir, "thread_count.json", {"thread_count": "1000"})
|
|
self._write_json(temp_dir, "process_count.json", {"process_count": "60"})
|
|
self._write_json(temp_dir, "node_thread_counts.json", {"mainland-worker-01": 1200})
|
|
self._write_json(temp_dir, "node_process_counts.json", {"mainland-worker-01": 60})
|
|
|
|
previous_cwd = os.getcwd()
|
|
try:
|
|
os.chdir(temp_dir)
|
|
limits = _resolve_db_pool_limits()
|
|
finally:
|
|
os.chdir(previous_cwd)
|
|
|
|
self.assertEqual(4, limits["pool_size"])
|
|
self.assertEqual(1, limits["pool_warm_size"])
|
|
self.assertEqual(2, limits["pool_idle_keep_max"])
|
|
self.assertEqual(60, limits["scaling_hints"]["process_count"])
|
|
|
|
def test_explicit_pool_env_overrides_win_over_auto_scaling(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir, patch.object(
|
|
config, "NODE_CODE", "mainland-controller-01"
|
|
), patch.dict(
|
|
os.environ,
|
|
{
|
|
"DB_POOL_SIZE": "16",
|
|
"DB_POOL_WARM_SIZE": "4",
|
|
"DB_POOL_IDLE_KEEP_MAX": "8",
|
|
},
|
|
clear=False,
|
|
):
|
|
self._write_json(temp_dir, "thread_count.json", {"thread_count": "1000"})
|
|
self._write_json(temp_dir, "process_count.json", {"process_count": "80"})
|
|
|
|
previous_cwd = os.getcwd()
|
|
try:
|
|
os.chdir(temp_dir)
|
|
limits = _resolve_db_pool_limits()
|
|
finally:
|
|
os.chdir(previous_cwd)
|
|
|
|
self.assertEqual(16, limits["pool_size"])
|
|
self.assertEqual(4, limits["pool_warm_size"])
|
|
self.assertEqual(8, limits["pool_idle_keep_max"])
|
|
|
|
def test_pool_limits_can_resolve_parent_process_override_from_absolute_config_root(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir, patch.object(
|
|
config, "NODE_CODE", "mainland-controller-01-u"
|
|
), patch.dict(
|
|
os.environ,
|
|
{
|
|
"WORKER_PARENT_NODE_CODE": "mainland-controller-01",
|
|
"DB_POOL_SIZE": "",
|
|
"DB_POOL_WARM_SIZE": "",
|
|
"DB_POOL_IDLE_KEEP_MAX": "",
|
|
"DOMAINCHECK_DB_POOL_TOTAL_BUDGET": "",
|
|
},
|
|
clear=False,
|
|
):
|
|
self._write_json(temp_dir, "thread_count.json", {"thread_count": "1000"})
|
|
self._write_json(temp_dir, "process_count.json", {"process_count": "80"})
|
|
self._write_json(temp_dir, "node_thread_counts.json", {"mainland-controller-01": 1000})
|
|
self._write_json(temp_dir, "node_process_counts.json", {"mainland-controller-01": 80})
|
|
|
|
with patch(
|
|
"app.utils.database._local_config_search_roots",
|
|
return_value=[Path("/nonexistent/domainCheck"), Path(temp_dir)],
|
|
):
|
|
limits = _resolve_db_pool_limits()
|
|
|
|
self.assertEqual(80, limits["scaling_hints"]["process_count"])
|
|
self.assertEqual("mainland-controller-01", limits["scaling_hints"]["parent_node_code"])
|
|
self.assertEqual(9, limits["pool_size"])
|
|
|
|
def test_controller_high_process_pool_keeps_ten_connections_at_hundred_processes(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir, patch.object(
|
|
config, "NODE_CODE", "mainland-controller-01-aa"
|
|
), patch.dict(
|
|
os.environ,
|
|
{
|
|
"WORKER_PARENT_NODE_CODE": "mainland-controller-01",
|
|
"DB_POOL_SIZE": "",
|
|
"DB_POOL_WARM_SIZE": "",
|
|
"DB_POOL_IDLE_KEEP_MAX": "",
|
|
"DOMAINCHECK_DB_POOL_TOTAL_BUDGET": "",
|
|
},
|
|
clear=False,
|
|
):
|
|
self._write_json(temp_dir, "thread_count.json", {"thread_count": "1000"})
|
|
self._write_json(temp_dir, "process_count.json", {"process_count": "100"})
|
|
self._write_json(temp_dir, "node_thread_counts.json", {"mainland-controller-01": 1000})
|
|
self._write_json(temp_dir, "node_process_counts.json", {"mainland-controller-01": 100})
|
|
|
|
previous_cwd = os.getcwd()
|
|
try:
|
|
os.chdir(temp_dir)
|
|
limits = _resolve_db_pool_limits()
|
|
finally:
|
|
os.chdir(previous_cwd)
|
|
|
|
self.assertEqual(100, limits["scaling_hints"]["process_count"])
|
|
self.assertEqual(10, limits["pool_size"])
|
|
self.assertEqual(2, limits["pool_warm_size"])
|
|
self.assertEqual(5, limits["pool_idle_keep_max"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|