add runtime adapters, scheduler, reconciler, and their unit tests
This commit is contained in:
parent
d1976a5fd8
commit
b63d69c81d
10 changed files with 1471 additions and 28 deletions
|
|
@ -1 +1,286 @@
|
|||
"""EC2 runtime unit tests — Plan 02."""
|
||||
"""Unit tests for the EC2 runtime adapter using botocore Stubber."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import patch
|
||||
|
||||
import boto3
|
||||
import pytest
|
||||
from botocore.stub import Stubber
|
||||
|
||||
from nix_builder_autoscaler.config import AwsConfig
|
||||
from nix_builder_autoscaler.runtime.base import RuntimeError as RuntimeAdapterError
|
||||
from nix_builder_autoscaler.runtime.ec2 import EC2Runtime
|
||||
|
||||
|
||||
def _make_config():
|
||||
return AwsConfig(
|
||||
region="us-east-1",
|
||||
launch_template_id="lt-abc123",
|
||||
subnet_ids=["subnet-aaa", "subnet-bbb"],
|
||||
security_group_ids=["sg-111"],
|
||||
instance_profile_arn="arn:aws:iam::123456789012:instance-profile/nix-builder",
|
||||
)
|
||||
|
||||
|
||||
def _make_runtime(stubber, ec2_client, **kwargs):
|
||||
config = kwargs.pop("config", _make_config())
|
||||
environment = kwargs.pop("environment", "dev")
|
||||
stubber.activate()
|
||||
return EC2Runtime(config, environment=environment, _client=ec2_client)
|
||||
|
||||
|
||||
class TestLaunchSpot:
|
||||
def test_correct_params_and_returns_instance_id(self):
|
||||
config = _make_config()
|
||||
ec2_client = boto3.client("ec2", region_name="us-east-1")
|
||||
stubber = Stubber(ec2_client)
|
||||
|
||||
expected_params = {
|
||||
"MinCount": 1,
|
||||
"MaxCount": 1,
|
||||
"LaunchTemplate": {
|
||||
"LaunchTemplateId": "lt-abc123",
|
||||
"Version": "$Latest",
|
||||
},
|
||||
"InstanceMarketOptions": {
|
||||
"MarketType": "spot",
|
||||
"SpotOptions": {
|
||||
"SpotInstanceType": "one-time",
|
||||
"InstanceInterruptionBehavior": "terminate",
|
||||
},
|
||||
},
|
||||
"SubnetId": "subnet-aaa",
|
||||
"UserData": "#!/bin/bash\necho hello",
|
||||
"TagSpecifications": [
|
||||
{
|
||||
"ResourceType": "instance",
|
||||
"Tags": [
|
||||
{"Key": "Name", "Value": "nix-builder-slot001"},
|
||||
{"Key": "AutoscalerSlot", "Value": "slot001"},
|
||||
{"Key": "ManagedBy", "Value": "nix-builder-autoscaler"},
|
||||
{"Key": "Service", "Value": "nix-builder"},
|
||||
{"Key": "Environment", "Value": "dev"},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
response = {
|
||||
"Instances": [{"InstanceId": "i-12345678"}],
|
||||
"OwnerId": "123456789012",
|
||||
}
|
||||
|
||||
stubber.add_response("run_instances", response, expected_params)
|
||||
runtime = _make_runtime(stubber, ec2_client, config=config)
|
||||
|
||||
iid = runtime.launch_spot("slot001", "#!/bin/bash\necho hello")
|
||||
assert iid == "i-12345678"
|
||||
stubber.assert_no_pending_responses()
|
||||
|
||||
def test_round_robin_subnets(self):
|
||||
config = _make_config()
|
||||
ec2_client = boto3.client("ec2", region_name="us-east-1")
|
||||
stubber = Stubber(ec2_client)
|
||||
|
||||
# Two launches should use subnet-aaa then subnet-bbb
|
||||
for _subnet in ["subnet-aaa", "subnet-bbb"]:
|
||||
stubber.add_response(
|
||||
"run_instances",
|
||||
{"Instances": [{"InstanceId": "i-abc"}], "OwnerId": "123"},
|
||||
)
|
||||
|
||||
runtime = _make_runtime(stubber, ec2_client, config=config)
|
||||
runtime.launch_spot("slot001", "")
|
||||
runtime.launch_spot("slot002", "")
|
||||
stubber.assert_no_pending_responses()
|
||||
|
||||
|
||||
class TestDescribeInstance:
|
||||
def test_normalizes_response(self):
|
||||
ec2_client = boto3.client("ec2", region_name="us-east-1")
|
||||
stubber = Stubber(ec2_client)
|
||||
|
||||
launch_time = datetime(2026, 1, 15, 12, 30, 0, tzinfo=UTC)
|
||||
response = {
|
||||
"Reservations": [
|
||||
{
|
||||
"Instances": [
|
||||
{
|
||||
"InstanceId": "i-running1",
|
||||
"State": {"Code": 16, "Name": "running"},
|
||||
"LaunchTime": launch_time,
|
||||
"Tags": [
|
||||
{"Key": "AutoscalerSlot", "Value": "slot001"},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
stubber.add_response(
|
||||
"describe_instances",
|
||||
response,
|
||||
{"InstanceIds": ["i-running1"]},
|
||||
)
|
||||
runtime = _make_runtime(stubber, ec2_client)
|
||||
|
||||
info = runtime.describe_instance("i-running1")
|
||||
assert info["state"] == "running"
|
||||
assert info["tailscale_ip"] is None
|
||||
assert info["launch_time"] == launch_time.isoformat()
|
||||
|
||||
def test_missing_instance_returns_terminated(self):
|
||||
ec2_client = boto3.client("ec2", region_name="us-east-1")
|
||||
stubber = Stubber(ec2_client)
|
||||
|
||||
stubber.add_response(
|
||||
"describe_instances",
|
||||
{"Reservations": []},
|
||||
{"InstanceIds": ["i-gone"]},
|
||||
)
|
||||
runtime = _make_runtime(stubber, ec2_client)
|
||||
|
||||
info = runtime.describe_instance("i-gone")
|
||||
assert info["state"] == "terminated"
|
||||
assert info["tailscale_ip"] is None
|
||||
assert info["launch_time"] is None
|
||||
|
||||
|
||||
class TestListManagedInstances:
|
||||
def test_filters_by_tag(self):
|
||||
ec2_client = boto3.client("ec2", region_name="us-east-1")
|
||||
stubber = Stubber(ec2_client)
|
||||
|
||||
expected_params = {
|
||||
"Filters": [
|
||||
{"Name": "tag:ManagedBy", "Values": ["nix-builder-autoscaler"]},
|
||||
{
|
||||
"Name": "instance-state-name",
|
||||
"Values": ["pending", "running", "shutting-down", "stopping"],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
response = {
|
||||
"Reservations": [
|
||||
{
|
||||
"Instances": [
|
||||
{
|
||||
"InstanceId": "i-aaa",
|
||||
"State": {"Code": 16, "Name": "running"},
|
||||
"Tags": [
|
||||
{"Key": "AutoscalerSlot", "Value": "slot001"},
|
||||
{"Key": "ManagedBy", "Value": "nix-builder-autoscaler"},
|
||||
],
|
||||
},
|
||||
{
|
||||
"InstanceId": "i-bbb",
|
||||
"State": {"Code": 0, "Name": "pending"},
|
||||
"Tags": [
|
||||
{"Key": "AutoscalerSlot", "Value": "slot002"},
|
||||
{"Key": "ManagedBy", "Value": "nix-builder-autoscaler"},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
stubber.add_response("describe_instances", response, expected_params)
|
||||
runtime = _make_runtime(stubber, ec2_client)
|
||||
|
||||
managed = runtime.list_managed_instances()
|
||||
assert len(managed) == 2
|
||||
assert managed[0]["instance_id"] == "i-aaa"
|
||||
assert managed[0]["state"] == "running"
|
||||
assert managed[0]["slot_id"] == "slot001"
|
||||
assert managed[1]["instance_id"] == "i-bbb"
|
||||
assert managed[1]["state"] == "pending"
|
||||
assert managed[1]["slot_id"] == "slot002"
|
||||
|
||||
|
||||
class TestTerminateInstance:
|
||||
def test_calls_terminate_api(self):
|
||||
ec2_client = boto3.client("ec2", region_name="us-east-1")
|
||||
stubber = Stubber(ec2_client)
|
||||
|
||||
response = {
|
||||
"TerminatingInstances": [
|
||||
{
|
||||
"InstanceId": "i-kill",
|
||||
"CurrentState": {"Code": 32, "Name": "shutting-down"},
|
||||
"PreviousState": {"Code": 16, "Name": "running"},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
stubber.add_response(
|
||||
"terminate_instances",
|
||||
response,
|
||||
{"InstanceIds": ["i-kill"]},
|
||||
)
|
||||
runtime = _make_runtime(stubber, ec2_client)
|
||||
|
||||
# Should not raise
|
||||
runtime.terminate_instance("i-kill")
|
||||
stubber.assert_no_pending_responses()
|
||||
|
||||
|
||||
class TestErrorClassification:
|
||||
def test_insufficient_capacity_classified(self):
|
||||
ec2_client = boto3.client("ec2", region_name="us-east-1")
|
||||
stubber = Stubber(ec2_client)
|
||||
|
||||
stubber.add_client_error(
|
||||
"run_instances",
|
||||
service_error_code="InsufficientInstanceCapacity",
|
||||
service_message="Insufficient capacity",
|
||||
)
|
||||
runtime = _make_runtime(stubber, ec2_client)
|
||||
|
||||
with pytest.raises(RuntimeAdapterError) as exc_info:
|
||||
runtime.launch_spot("slot001", "#!/bin/bash")
|
||||
assert exc_info.value.category == "capacity_unavailable"
|
||||
|
||||
@patch("nix_builder_autoscaler.runtime.ec2.time.sleep")
|
||||
def test_request_limit_exceeded_retried(self, mock_sleep):
|
||||
ec2_client = boto3.client("ec2", region_name="us-east-1")
|
||||
stubber = Stubber(ec2_client)
|
||||
|
||||
# First call: throttled
|
||||
stubber.add_client_error(
|
||||
"run_instances",
|
||||
service_error_code="RequestLimitExceeded",
|
||||
service_message="Rate exceeded",
|
||||
)
|
||||
# Second call: success
|
||||
stubber.add_response(
|
||||
"run_instances",
|
||||
{"Instances": [{"InstanceId": "i-retry123"}], "OwnerId": "123"},
|
||||
)
|
||||
runtime = _make_runtime(stubber, ec2_client)
|
||||
|
||||
iid = runtime.launch_spot("slot001", "#!/bin/bash")
|
||||
assert iid == "i-retry123"
|
||||
assert mock_sleep.called
|
||||
stubber.assert_no_pending_responses()
|
||||
|
||||
@patch("nix_builder_autoscaler.runtime.ec2.time.sleep")
|
||||
def test_request_limit_exceeded_exhausted(self, mock_sleep):
|
||||
"""After max retries, RequestLimitExceeded raises with 'throttled' category."""
|
||||
ec2_client = boto3.client("ec2", region_name="us-east-1")
|
||||
stubber = Stubber(ec2_client)
|
||||
|
||||
# 4 errors (max_retries=3: attempt 0,1,2,3 all fail)
|
||||
for _ in range(4):
|
||||
stubber.add_client_error(
|
||||
"run_instances",
|
||||
service_error_code="RequestLimitExceeded",
|
||||
service_message="Rate exceeded",
|
||||
)
|
||||
runtime = _make_runtime(stubber, ec2_client)
|
||||
|
||||
with pytest.raises(RuntimeAdapterError) as exc_info:
|
||||
runtime.launch_spot("slot001", "#!/bin/bash")
|
||||
assert exc_info.value.category == "throttled"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue