mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-11-26 17:41:59 +00:00
146 lines
4.9 KiB
Python
146 lines
4.9 KiB
Python
#!/usr/bin/env python
|
|
|
|
import unittest
|
|
from dataclasses import dataclass
|
|
from typing import Any, List
|
|
|
|
from autoscale_runners_lambda.app import set_capacity, Queue
|
|
|
|
|
|
@dataclass
|
|
class TestCase:
|
|
name: str
|
|
min_size: int
|
|
desired_capacity: int
|
|
max_size: int
|
|
queues: List[Queue]
|
|
expected_capacity: int
|
|
|
|
|
|
class TestSetCapacity(unittest.TestCase):
|
|
class FakeClient:
|
|
def __init__(self):
|
|
self._expected_data = {} # type: dict
|
|
self._expected_capacity = -1
|
|
|
|
@property
|
|
def expected_data(self) -> dict:
|
|
"""a one-time property"""
|
|
data, self._expected_data = self._expected_data, {}
|
|
return data
|
|
|
|
@expected_data.setter
|
|
def expected_data(self, value: dict) -> None:
|
|
self._expected_data = value
|
|
|
|
@property
|
|
def expected_capacity(self) -> int:
|
|
"""one-time property"""
|
|
capacity, self._expected_capacity = self._expected_capacity, -1
|
|
return capacity
|
|
|
|
def describe_auto_scaling_groups(self, **kwargs: Any) -> dict:
|
|
_ = kwargs
|
|
return self.expected_data
|
|
|
|
def set_desired_capacity(self, **kwargs: Any) -> None:
|
|
self._expected_capacity = kwargs["DesiredCapacity"]
|
|
|
|
def data_helper(
|
|
self, name: str, min_size: int, desired_capacity: int, max_size: int
|
|
) -> None:
|
|
self.expected_data = {
|
|
"AutoScalingGroups": [
|
|
{
|
|
"AutoScalingGroupName": name,
|
|
"DesiredCapacity": desired_capacity,
|
|
"MinSize": min_size,
|
|
"MaxSize": max_size,
|
|
}
|
|
]
|
|
}
|
|
|
|
def setUp(self):
|
|
self.client = self.FakeClient()
|
|
|
|
def test_normal_cases(self):
|
|
test_cases = (
|
|
# Do not change capacity
|
|
TestCase("noqueue", 1, 13, 20, [Queue("in_progress", 155, "noqueue")], -1),
|
|
TestCase("w/reserve", 1, 13, 20, [Queue("queued", 17, "w/reserve")], -1),
|
|
# Increase capacity
|
|
TestCase("increase", 1, 13, 20, [Queue("queued", 23, "increase")], 15),
|
|
TestCase(
|
|
"style-checker", 1, 13, 20, [Queue("queued", 33, "style-checker")], 15
|
|
),
|
|
TestCase("increase", 1, 13, 20, [Queue("queued", 18, "increase")], 14),
|
|
TestCase("increase", 1, 13, 20, [Queue("queued", 183, "increase")], 20),
|
|
TestCase(
|
|
"increase-w/o reserve",
|
|
1,
|
|
13,
|
|
20,
|
|
[
|
|
Queue("in_progress", 11, "increase-w/o reserve"),
|
|
Queue("queued", 12, "increase-w/o reserve"),
|
|
],
|
|
15,
|
|
),
|
|
TestCase("lower-min", 10, 5, 20, [Queue("queued", 5, "lower-min")], 10),
|
|
# Decrease capacity
|
|
TestCase("w/reserve", 1, 13, 20, [Queue("queued", 5, "w/reserve")], 9),
|
|
TestCase(
|
|
"style-checker", 1, 13, 20, [Queue("queued", 5, "style-checker")], 5
|
|
),
|
|
TestCase("w/reserve", 1, 23, 20, [Queue("queued", 17, "w/reserve")], 20),
|
|
TestCase("decrease", 1, 13, 20, [Queue("in_progress", 3, "decrease")], 8),
|
|
TestCase(
|
|
"style-checker",
|
|
1,
|
|
13,
|
|
20,
|
|
[Queue("in_progress", 5, "style-checker")],
|
|
5,
|
|
),
|
|
)
|
|
for t in test_cases:
|
|
self.client.data_helper(t.name, t.min_size, t.desired_capacity, t.max_size)
|
|
set_capacity(t.name, t.queues, self.client, False)
|
|
self.assertEqual(t.expected_capacity, self.client.expected_capacity, t.name)
|
|
|
|
def test_exceptions(self):
|
|
test_cases = (
|
|
(
|
|
TestCase(
|
|
"different names",
|
|
1,
|
|
1,
|
|
1,
|
|
[Queue("queued", 5, "another name")],
|
|
-1,
|
|
),
|
|
AssertionError,
|
|
),
|
|
(TestCase("wrong queue len", 1, 1, 1, [], -1), AssertionError),
|
|
(
|
|
TestCase(
|
|
"wrong queue", 1, 1, 1, [Queue("wrong", 1, "wrong queue")], -1 # type: ignore
|
|
),
|
|
ValueError,
|
|
),
|
|
)
|
|
for t, error in test_cases:
|
|
with self.assertRaises(error):
|
|
self.client.data_helper(
|
|
t.name, t.min_size, t.desired_capacity, t.max_size
|
|
)
|
|
set_capacity(t.name, t.queues, self.client, False)
|
|
|
|
with self.assertRaises(AssertionError):
|
|
self.client.expected_data = {"AutoScalingGroups": [1, 2]}
|
|
set_capacity(
|
|
"wrong number of ASGs",
|
|
[Queue("queued", 1, "wrong number of ASGs")],
|
|
self.client,
|
|
)
|