diff --git a/cloudproxy/providers/aws/main.py b/cloudproxy/providers/aws/main.py index 8c19e07..53c60d5 100644 --- a/cloudproxy/providers/aws/main.py +++ b/cloudproxy/providers/aws/main.py @@ -67,8 +67,8 @@ def aws_check_alive(instance_config=None): "default" ) - ip_ready = [] - pending_ips = [] + ip_ready = set() + pending_ips = set() instances_to_recycle = [] # First pass: identify healthy and pending instances @@ -78,10 +78,7 @@ def aws_check_alive(instance_config=None): datetime.timezone.utc ) - instance["Instances"][0]["LaunchTime"] - if config["age_limit"] > 0 and elapsed > datetime.timedelta(seconds=config["age_limit"]): - # Queue for potential recycling - instances_to_recycle.append((instance, elapsed)) - elif instance["Instances"][0]["State"]["Name"] == "stopped": + if instance["Instances"][0]["State"]["Name"] == "stopped": logger.info( f"Waking up: AWS {instance_config.get('display_name', 'default')} -> Instance " + instance["Instances"][0]["InstanceId"] ) @@ -99,13 +96,16 @@ def aws_check_alive(instance_config=None): f"Pending: AWS {instance_config.get('display_name', 'default')} -> " + instance["Instances"][0]["PublicIpAddress"] ) if "PublicIpAddress" in instance["Instances"][0]: - pending_ips.append(instance["Instances"][0]["PublicIpAddress"]) + pending_ips.add(instance["Instances"][0]["PublicIpAddress"]) # Must be "running" if none of the above, check if alive or not. elif check_alive(instance["Instances"][0]["PublicIpAddress"]): logger.info( f"Alive: AWS {instance_config.get('display_name', 'default')} -> " + instance["Instances"][0]["PublicIpAddress"] ) - ip_ready.append(instance["Instances"][0]["PublicIpAddress"]) + ip_ready.add(instance["Instances"][0]["PublicIpAddress"]) + if config["age_limit"] > 0 and elapsed > datetime.timedelta(seconds=config["age_limit"]): + # Queue for potential recycling + instances_to_recycle.append((instance, elapsed)) else: if elapsed > datetime.timedelta(minutes=10): delete_proxy(instance["Instances"][0]["InstanceId"], instance_config) @@ -118,7 +118,7 @@ def aws_check_alive(instance_config=None): f"Waiting: AWS {instance_config.get('display_name', 'default')} -> " + instance["Instances"][0]["PublicIpAddress"] ) if "PublicIpAddress" in instance["Instances"][0]: - pending_ips.append(instance["Instances"][0]["PublicIpAddress"]) + pending_ips.add(instance["Instances"][0]["PublicIpAddress"]) except (TypeError, KeyError): logger.info(f"Pending: AWS {instance_config.get('display_name', 'default')} -> allocating ip") @@ -147,6 +147,7 @@ def aws_check_alive(instance_config=None): # Mark as recycling and delete rolling_manager.mark_proxy_recycling("aws", instance_name, instance_ip) delete_proxy(inst["Instances"][0]["InstanceId"], instance_config) + ip_ready.discard(instance_ip) rolling_manager.mark_proxy_recycled("aws", instance_name, instance_ip) logger.info( f"Rolling deployment: Recycled AWS {instance_config.get('display_name', 'default')} instance (age limit) -> {instance_ip}" @@ -160,6 +161,7 @@ def aws_check_alive(instance_config=None): for inst, elapsed in instances_to_recycle: delete_proxy(inst["Instances"][0]["InstanceId"], instance_config) if "PublicIpAddress" in inst["Instances"][0]: + ip_ready.discard(inst["Instances"][0]["PublicIpAddress"]) logger.info( f"Recycling AWS {instance_config.get('display_name', 'default')} instance, reached age limit -> " + inst["Instances"][0]["PublicIpAddress"] ) @@ -168,7 +170,7 @@ def aws_check_alive(instance_config=None): f"Recycling AWS {instance_config.get('display_name', 'default')} instance, reached age limit -> " + inst["Instances"][0]["InstanceId"] ) - return ip_ready + return list(ip_ready) def aws_check_delete(instance_config=None): diff --git a/cloudproxy/providers/digitalocean/main.py b/cloudproxy/providers/digitalocean/main.py index 621c7b4..14623c4 100644 --- a/cloudproxy/providers/digitalocean/main.py +++ b/cloudproxy/providers/digitalocean/main.py @@ -69,8 +69,8 @@ def do_check_alive(instance_config=None): "default" ) - ip_ready = [] - pending_ips = [] + ip_ready = set() + pending_ips = set() droplets_to_recycle = [] # First pass: identify healthy and pending droplets @@ -81,18 +81,18 @@ def do_check_alive(instance_config=None): if created_at is None: # If parsing fails but doesn't raise an exception, log and continue logger.info(f"Pending: DO {display_name} allocating (invalid timestamp)") - pending_ips.append(str(droplet.ip_address)) + pending_ips.add(str(droplet.ip_address)) continue # Calculate elapsed time elapsed = datetime.datetime.now(datetime.timezone.utc) - created_at - # Check if the droplet has reached the age limit - if config["age_limit"] > 0 and elapsed > datetime.timedelta(seconds=config["age_limit"]): - droplets_to_recycle.append((droplet, elapsed)) - elif check_alive(droplet.ip_address): + if check_alive(droplet.ip_address): logger.info(f"Alive: DO {display_name} -> {str(droplet.ip_address)}") - ip_ready.append(droplet.ip_address) + ip_ready.add(droplet.ip_address) + # Check if the droplet has reached the age limit + if config["age_limit"] > 0 and elapsed > datetime.timedelta(seconds=config["age_limit"]): + droplets_to_recycle.append((droplet, elapsed)) else: # Check if the droplet has been pending for too long if elapsed > datetime.timedelta(minutes=10): @@ -102,12 +102,12 @@ def do_check_alive(instance_config=None): ) else: logger.info(f"Waiting: DO {display_name} -> {str(droplet.ip_address)}") - pending_ips.append(str(droplet.ip_address)) + pending_ips.add(str(droplet.ip_address)) except TypeError: # This happens when dateparser.parse raises a TypeError logger.info(f"Pending: DO {display_name} allocating") if hasattr(droplet, 'ip_address'): - pending_ips.append(str(droplet.ip_address)) + pending_ips.add(str(droplet.ip_address)) # Update rolling manager with current proxy health status rolling_manager.update_proxy_health("digitalocean", instance_name, ip_ready, pending_ips) @@ -133,6 +133,7 @@ def do_check_alive(instance_config=None): # Mark as recycling and delete rolling_manager.mark_proxy_recycling("digitalocean", instance_name, droplet_ip) delete_proxy(droplet, instance_config) + ip_ready.discard(droplet.ip_address) rolling_manager.mark_proxy_recycled("digitalocean", instance_name, droplet_ip) logger.info( f"Rolling deployment: Recycled DO {display_name} droplet (age limit) -> {droplet_ip}" @@ -145,11 +146,12 @@ def do_check_alive(instance_config=None): # Standard non-rolling recycling for droplet, elapsed in droplets_to_recycle: delete_proxy(droplet, instance_config) + ip_ready.discard(droplet.ip_address) logger.info( f"Recycling DO {display_name} droplet, reached age limit -> {str(droplet.ip_address)}" ) - return ip_ready + return list(ip_ready) def do_check_delete(instance_config=None): diff --git a/cloudproxy/providers/gcp/main.py b/cloudproxy/providers/gcp/main.py index 8390b6c..b4ba9f7 100644 --- a/cloudproxy/providers/gcp/main.py +++ b/cloudproxy/providers/gcp/main.py @@ -62,8 +62,8 @@ def gcp_check_alive(instance_config=None): "default" ) - ip_ready = [] - pending_ips = [] + ip_ready = set() + pending_ips = set() instances_to_recycle = [] for instance in list_instances(instance_config): @@ -72,11 +72,7 @@ def gcp_check_alive(instance_config=None): datetime.timezone.utc ) - datetime.datetime.strptime(instance["creationTimestamp"], '%Y-%m-%dT%H:%M:%S.%f%z') - if config["age_limit"] > 0 and elapsed > datetime.timedelta(seconds=config["age_limit"]): - # Queue for potential recycling - instances_to_recycle.append((instance, elapsed)) - - elif instance['status'] == "TERMINATED": + if instance['status'] == "TERMINATED": logger.info("Waking up: GCP -> Instance " + instance['name']) started = start_proxy(instance['name'], instance_config) if not started: @@ -92,14 +88,17 @@ def gcp_check_alive(instance_config=None): msg = f"{instance['name']} {access_configs['natIP'] if 'natIP' in access_configs else ''}" logger.info("Provisioning: GCP -> " + msg) if 'natIP' in access_configs: - pending_ips.append(access_configs['natIP']) + pending_ips.add(access_configs['natIP']) # If none of the above, check if alive or not. elif check_alive(instance['networkInterfaces'][0]['accessConfigs'][0]['natIP']): access_configs = instance['networkInterfaces'][0]['accessConfigs'][0] msg = f"{instance['name']} {access_configs['natIP']}" logger.info("Alive: GCP -> " + msg) - ip_ready.append(access_configs['natIP']) + ip_ready.add(access_configs['natIP']) + if config["age_limit"] > 0 and elapsed > datetime.timedelta(seconds=config["age_limit"]): + # Queue for potential recycling + instances_to_recycle.append((instance, elapsed)) else: access_configs = instance['networkInterfaces'][0]['accessConfigs'][0] @@ -110,7 +109,7 @@ def gcp_check_alive(instance_config=None): else: logger.info("Waiting: GCP -> " + msg) if 'natIP' in access_configs: - pending_ips.append(access_configs['natIP']) + pending_ips.add(access_configs['natIP']) except (TypeError, KeyError): logger.info("Pending: GCP -> Allocating IP") @@ -140,6 +139,7 @@ def gcp_check_alive(instance_config=None): # Mark as recycling and delete rolling_manager.mark_proxy_recycling("gcp", instance_name, instance_ip) delete_proxy(inst['name'], instance_config) + ip_ready.discard(instance_ip) rolling_manager.mark_proxy_recycled("gcp", instance_name, instance_ip) logger.info(f"Rolling deployment: Recycled GCP instance (age limit) -> {inst['name']} {instance_ip}") else: @@ -150,9 +150,10 @@ def gcp_check_alive(instance_config=None): access_configs = inst['networkInterfaces'][0]['accessConfigs'][0] msg = f"{inst['name']} {access_configs['natIP'] if 'natIP' in access_configs else ''}" delete_proxy(inst['name'], instance_config) + ip_ready.discard(access_configs['natIP']) logger.info("Recycling instance, reached age limit -> " + msg) - return ip_ready + return list(ip_ready) def gcp_check_delete(instance_config=None): """ diff --git a/cloudproxy/providers/hetzner/main.py b/cloudproxy/providers/hetzner/main.py index 77b9a8e..ba3121b 100644 --- a/cloudproxy/providers/hetzner/main.py +++ b/cloudproxy/providers/hetzner/main.py @@ -66,20 +66,20 @@ def hetzner_check_alive(instance_config=None): "default" ) - ip_ready = [] - pending_ips = [] + ip_ready = set() + pending_ips = set() proxies_to_recycle = [] for proxy in list_proxies(instance_config): elapsed = datetime.datetime.now( datetime.timezone.utc ) - dateparser.parse(str(proxy.created)) - if config["age_limit"] > 0 and elapsed > datetime.timedelta(seconds=config["age_limit"]): - # Queue for potential recycling - proxies_to_recycle.append((proxy, elapsed)) - elif check_alive(proxy.public_net.ipv4.ip): + if check_alive(proxy.public_net.ipv4.ip): logger.info(f"Alive: Hetzner {display_name} -> {str(proxy.public_net.ipv4.ip)}") - ip_ready.append(proxy.public_net.ipv4.ip) + ip_ready.add(proxy.public_net.ipv4.ip) + if config["age_limit"] > 0 and elapsed > datetime.timedelta(seconds=config["age_limit"]): + # Queue for potential recycling + proxies_to_recycle.append((proxy, elapsed)) else: if elapsed > datetime.timedelta(minutes=10): delete_proxy(proxy, instance_config) @@ -88,7 +88,7 @@ def hetzner_check_alive(instance_config=None): ) else: logger.info(f"Waiting: Hetzner {display_name} -> {str(proxy.public_net.ipv4.ip)}") - pending_ips.append(str(proxy.public_net.ipv4.ip)) + pending_ips.add(str(proxy.public_net.ipv4.ip)) # Update rolling manager with current proxy health status rolling_manager.update_proxy_health("hetzner", instance_name, ip_ready, pending_ips) @@ -114,6 +114,7 @@ def hetzner_check_alive(instance_config=None): # Mark as recycling and delete rolling_manager.mark_proxy_recycling("hetzner", instance_name, proxy_ip) delete_proxy(prox, instance_config) + ip_ready.discard(proxy_ip) rolling_manager.mark_proxy_recycled("hetzner", instance_name, proxy_ip) logger.info(f"Rolling deployment: Recycled Hetzner {display_name} proxy (age limit) -> {proxy_ip}") else: @@ -122,9 +123,10 @@ def hetzner_check_alive(instance_config=None): # Standard non-rolling recycling for prox, elapsed in proxies_to_recycle: delete_proxy(prox, instance_config) + ip_ready.discard(prox.public_net.ipv4.ip) logger.info(f"Recycling Hetzner {display_name} proxy, reached age limit -> {str(prox.public_net.ipv4.ip)}") - return ip_ready + return list(ip_ready) def hetzner_check_delete(instance_config=None): diff --git a/cloudproxy/providers/rolling.py b/cloudproxy/providers/rolling.py index a25fd0d..e57d8ae 100644 --- a/cloudproxy/providers/rolling.py +++ b/cloudproxy/providers/rolling.py @@ -149,8 +149,8 @@ def update_proxy_health( self, provider: str, instance: str, - healthy_ips: List[str], - pending_ips: List[str] = None + healthy_ips: Set[str], + pending_ips: Set[str] = None ): """ Update the health status of proxies for a provider instance. @@ -158,17 +158,17 @@ def update_proxy_health( Args: provider: The cloud provider name instance: The provider instance name - healthy_ips: List of IPs that are currently healthy - pending_ips: List of IPs that are pending (newly created) + healthy_ips: Set of IPs that are currently healthy + pending_ips: Set of IPs that are pending (newly created) """ state = self.get_state(provider, instance) # Update healthy proxies - state.healthy_proxies = set(healthy_ips) + state.healthy_proxies = healthy_ips # Update pending proxies if provided if pending_ips is not None: - state.pending = set(pending_ips) + state.pending = pending_ips # Clean up recycling list if proxies no longer exist existing_ips = state.healthy_proxies | state.pending diff --git a/cloudproxy/providers/vultr/main.py b/cloudproxy/providers/vultr/main.py index eca4de9..051c018 100644 --- a/cloudproxy/providers/vultr/main.py +++ b/cloudproxy/providers/vultr/main.py @@ -71,8 +71,8 @@ def vultr_check_alive(instance_config=None): "default" ) - ip_ready = [] - pending_ips = [] + ip_ready = set() + pending_ips = set() instances_to_recycle = [] for instance in list_instances(instance_config): @@ -89,15 +89,15 @@ def vultr_check_alive(instance_config=None): # Calculate elapsed time elapsed = datetime.datetime.now(datetime.timezone.utc) - created_at - # Check if the instance has reached the age limit - if config["age_limit"] > 0 and elapsed > datetime.timedelta( - seconds=config["age_limit"]): - # Queue for potential recycling - instances_to_recycle.append((instance, elapsed)) - elif instance.status == "active" and instance.ip_address and check_alive(instance.ip_address): + if instance.status == "active" and instance.ip_address and check_alive(instance.ip_address): logger.info( f"Alive: Vultr {display_name} -> {str(instance.ip_address)}") - ip_ready.append(instance.ip_address) + ip_ready.add(instance.ip_address) + # Check if the instance has reached the age limit + if config["age_limit"] > 0 and elapsed > datetime.timedelta( + seconds=config["age_limit"]): + # Queue for potential recycling + instances_to_recycle.append((instance, elapsed)) else: # Check if the instance has been pending for too long if elapsed > datetime.timedelta(minutes=10): @@ -109,12 +109,12 @@ def vultr_check_alive(instance_config=None): logger.info( f"Waiting: Vultr {display_name} -> {str(instance.ip_address)}") if instance.ip_address: - pending_ips.append(instance.ip_address) + pending_ips.add(instance.ip_address) except TypeError: # This happens when dateparser.parse raises a TypeError logger.info(f"Pending: Vultr {display_name} allocating") if hasattr(instance, 'ip_address') and instance.ip_address: - pending_ips.append(instance.ip_address) + pending_ips.add(instance.ip_address) # Update rolling manager with current proxy health status rolling_manager.update_proxy_health("vultr", instance_name, ip_ready, pending_ips) @@ -141,6 +141,7 @@ def vultr_check_alive(instance_config=None): # Mark as recycling and delete rolling_manager.mark_proxy_recycling("vultr", instance_name, instance_ip) delete_proxy(inst, instance_config) + ip_ready.discard(instance_ip) rolling_manager.mark_proxy_recycled("vultr", instance_name, instance_ip) logger.info( f"Rolling deployment: Recycled Vultr {display_name} instance (age limit) -> {instance_ip}" @@ -153,11 +154,12 @@ def vultr_check_alive(instance_config=None): # Standard non-rolling recycling for inst, elapsed in instances_to_recycle: delete_proxy(inst, instance_config) + ip_ready.discard(inst.ip_address) logger.info( f"Recycling Vultr {display_name} instance, reached age limit -> {str(inst.ip_address)}" ) - return ip_ready + return list(ip_ready) def vultr_check_delete(instance_config=None): diff --git a/tests/test_providers_aws_main.py b/tests/test_providers_aws_main.py index 636fc90..a2b592c 100644 --- a/tests/test_providers_aws_main.py +++ b/tests/test_providers_aws_main.py @@ -316,6 +316,55 @@ def test_aws_check_alive_age_limit_exceeded_directly(): config["age_limit"] = original_age_limit config["rolling_deployment"]["enabled"] = original_rolling +def test_aws_check_alive_age_limit_exceeded_with_rolling_deployment(): + """Test the aws_check_alive function with rolling deployment enabled and simulated old instance""" + # Save original age limit value + original_age_limit = config["age_limit"] + original_rolling = config["rolling_deployment"]["enabled"] + + try: + # Set age limit to a small value to make instances expire quickly + config["age_limit"] = 60 # 60 seconds + # Enable rolling deployment to allow for recycling + config["rolling_deployment"]["enabled"] = True + config["rolling_deployment"]["min_available"] = 1 + config["rolling_deployment"]["batch_size"] = 1 + + # Create a mock instance with a launch time far in the past + with patch('cloudproxy.providers.aws.main.list_instances') as mock_list_instances: + with patch('cloudproxy.providers.aws.main.check_alive') as mock_check_alive: + with patch('cloudproxy.providers.aws.main.delete_proxy') as mock_delete_proxy: + with patch('cloudproxy.providers.aws.main.start_proxy') as mock_start_proxy: + # Setup mocks with two old instances (from year 2000) + very_old_time = datetime.datetime(2000, 1, 1, tzinfo=timezone.utc) + mock_list_instances.return_value = [{ + "Instances": [{ + "InstanceId": "i-12345", + "PublicIpAddress": "1.2.3.4", + "State": {"Name": "running"}, + "LaunchTime": very_old_time + }]}, + {"Instances": [{ + "InstanceId": "i-67890", + "PublicIpAddress": "5.6.7.8", + "State": {"Name": "running"}, + "LaunchTime": very_old_time + }] + }] + mock_check_alive.return_value = True + mock_delete_proxy.return_value = True + + # Execute + result = aws_check_alive() + + # Verify + assert mock_delete_proxy.call_count == 1 # Should delete only one expired instance + assert len(result) == 1 # One IP in result as the instance was kept due to rolling deployment min_available rules + finally: + # Restore original settings + config["age_limit"] = original_age_limit + config["rolling_deployment"]["enabled"] = original_rolling + @patch('cloudproxy.providers.aws.main.list_instances') @patch('cloudproxy.providers.aws.main.delete_proxy') def test_aws_check_delete(mock_delete_proxy, mock_list_instances, setup_instances): diff --git a/tests/test_providers_digitalocean_main_coverage.py b/tests/test_providers_digitalocean_main_coverage.py index bf4375e..e8b6b1e 100644 --- a/tests/test_providers_digitalocean_main_coverage.py +++ b/tests/test_providers_digitalocean_main_coverage.py @@ -394,6 +394,50 @@ def test_do_check_alive_age_limit(mock_delete_proxy, mock_check_alive, mock_list config["age_limit"] = original_age_limit config["rolling_deployment"]["enabled"] = original_rolling +@patch('cloudproxy.providers.digitalocean.main.list_droplets') +@patch('cloudproxy.providers.digitalocean.main.check_alive') +@patch('cloudproxy.providers.digitalocean.main.delete_proxy') +def test_do_check_alive_age_limit_with_rolling_deployment(mock_delete_proxy, mock_check_alive, mock_list_droplets): + """Test age limit recycling of droplets with rolling deployment enabled""" + # Setup + # Setup mocks with two old instances (from year 2000) + very_old_time = datetime.datetime(2000, 1, 1, tzinfo=timezone.utc).isoformat() + + mock_droplets = [ + MockDroplet(1, "1.2.3.4", very_old_time), + MockDroplet(2, "5.6.7.8", very_old_time), + ] + mock_list_droplets.return_value = mock_droplets + mock_check_alive.return_value = True + mock_delete_proxy.return_value = True + + # Set age limit to 1 hour and enable rolling deployment + original_age_limit = config["age_limit"] + original_rolling = config["rolling_deployment"]["enabled"] + config["age_limit"] = 3600 # 1 hour in seconds + config["rolling_deployment"]["enabled"] = True + config["rolling_deployment"]["min_available"] = 1 + config["rolling_deployment"]["batch_size"] = 1 + + try: + # Execute + result = do_check_alive() + + # Verify - only the first droplet should be recycled due to rolling deployment min_available rules + assert len(result) == 1 + assert "1.2.3.4" not in result + assert "5.6.7.8" in result + assert mock_delete_proxy.call_count == 1 # Only one droplet should be deleted + + # Check that delete_proxy was called with the first droplet + # Don't verify the exact second parameter as implementation details might vary + args, _ = mock_delete_proxy.call_args + assert args[0] == mock_droplets[0] # First parameter should be the first droplet + finally: + # Restore original settings + config["age_limit"] = original_age_limit + config["rolling_deployment"]["enabled"] = original_rolling + @patch('cloudproxy.providers.digitalocean.main.list_droplets') @patch('cloudproxy.providers.digitalocean.main.check_alive') def test_do_check_alive_invalid_timestamp(mock_check_alive, mock_list_droplets, mock_droplets): diff --git a/tests/test_providers_gcp_main.py b/tests/test_providers_gcp_main.py index 97e3764..d3a14b5 100644 --- a/tests/test_providers_gcp_main.py +++ b/tests/test_providers_gcp_main.py @@ -270,6 +270,56 @@ def test_gcp_check_alive_age_limit_exceeded(mock_start_proxy, mock_delete_proxy, config["age_limit"] = original_age_limit config["rolling_deployment"]["enabled"] = original_rolling +@patch('cloudproxy.providers.gcp.main.check_alive') +@patch('cloudproxy.providers.gcp.main.list_instances') +@patch('cloudproxy.providers.gcp.main.delete_proxy') +@patch('cloudproxy.providers.gcp.main.start_proxy') +def test_gcp_check_alive_age_limit_exceeded_with_rolling_deployment(mock_start_proxy, mock_delete_proxy, mock_list_instances, mock_check_alive, setup_instances): + """Test checking alive for instances exceeding age limit with rolling deployment enabled""" + # Save original values + original_age_limit = config["age_limit"] + original_rolling = config["rolling_deployment"]["enabled"] + + try: + # Set age limit to a small value and enable rolling deployment + config["age_limit"] = 60 # 60 seconds + config["rolling_deployment"]["enabled"] = True + config["rolling_deployment"]["min_available"] = 1 + config["rolling_deployment"]["batch_size"] = 1 + + # Create two mock instance with a creation time far in the past + old_time = datetime.datetime.now(timezone.utc) - datetime.timedelta(seconds=120) + old_time_str = old_time.strftime('%Y-%m-%dT%H:%M:%S.%f%z') + old_instances = [ + { + "name": "old-instance-age-1", + "networkInterfaces": [{"accessConfigs": [{"natIP": "25.26.27.28"}]}], + "status": "RUNNING", + "creationTimestamp": old_time_str + }, + { + "name": "old-instance-age-2", + "networkInterfaces": [{"accessConfigs": [{"natIP": "29.30.31.32"}]}], + "status": "RUNNING", + "creationTimestamp": old_time_str + } + ] + mock_list_instances.return_value = old_instances + mock_check_alive.return_value = True # Instance is alive but old + mock_delete_proxy.return_value = True + + # Execute + result = gcp_check_alive() + + # Verify + assert mock_delete_proxy.call_count == 1 # Should delete only one instance + assert len(result) == 1 # One IP in result as the instance was kept due to rolling deployment min_available rules + finally: + # Restore original settings + config["age_limit"] = original_age_limit + config["rolling_deployment"]["enabled"] = original_rolling + + @patch('cloudproxy.providers.gcp.main.check_alive') @patch('cloudproxy.providers.gcp.main.list_instances') @patch('cloudproxy.providers.gcp.main.delete_proxy') diff --git a/tests/test_providers_hetzner_main.py b/tests/test_providers_hetzner_main.py index 99cfa4c..ba74312 100644 --- a/tests/test_providers_hetzner_main.py +++ b/tests/test_providers_hetzner_main.py @@ -109,11 +109,72 @@ def now(tz=None): mock_parse.return_value = datetime.datetime(2023, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc) # Run the function - hetzner_check_alive(mock_instance_config) + result = hetzner_check_alive(mock_instance_config) - # Use the instance config directly in assertion - mock_delete_proxy.assert_called_once_with(mock_proxy, mock_instance_config) - mock_check_alive.assert_not_called() + # Use the instance config directly in assertion + mock_delete_proxy.assert_called_once_with(mock_proxy, mock_instance_config) + assert len(result) == 0 + + @patch("cloudproxy.providers.hetzner.main.config", new_callable=MagicMock) + @patch("cloudproxy.providers.hetzner.main.check_alive") + @patch("cloudproxy.providers.hetzner.main.delete_proxy") + @patch("cloudproxy.providers.hetzner.main.list_proxies") + def test_hetzner_check_alive_recycling_with_rolling_deployment(self, mock_list_proxies, mock_delete_proxy, mock_check_alive, mock_config): + """Test recycling of Hetzner proxies based on age limit with rolling deployment.""" + # Setup proxy + mock_proxy = MagicMock(public_net=MagicMock(ipv4=MagicMock(ip="1.1.1.1"))) + mock_proxy.created = "2023-01-01T00:00:00Z" + mock_proxy_2 = MagicMock(public_net=MagicMock(ipv4=MagicMock(ip="2.2.2.2"))) + mock_proxy_2.created = "2023-01-01T00:00:00Z" + mock_list_proxies.return_value = [mock_proxy, mock_proxy_2] + mock_instance_config = {"display_name": "test", "scaling": {"min_scaling": 1}} + + # Configure the mock config + providers_dict = {"hetzner": {"instances": {"default": mock_instance_config}}} + + def config_getitem(key): + if key == "providers": + return providers_dict + elif key == "age_limit": + return 100 # Return an actual integer + elif key == "rolling_deployment": + return {"enabled": True, "min_available": 1, "batch_size": 1} + else: + return MagicMock() + + mock_config.__getitem__.side_effect = config_getitem + + # Create a module-level function to replace the datetime operations + # This function will be called when the implementation calculates elapsed time + def mock_elapsed_time(*args, **kwargs): + # Return a timedelta that's greater than age_limit (100 seconds) + return datetime.timedelta(seconds=101) + + # Patch the module-level function + with patch("cloudproxy.providers.hetzner.main.datetime") as mock_datetime: + # Setup the datetime mock to return our custom elapsed time + mock_datetime.timedelta = datetime.timedelta + mock_datetime.timezone = datetime.timezone + + # Create a mock datetime class with a now method + class MockDatetime: + @staticmethod + def now(tz=None): + return datetime.datetime(2023, 1, 1, 0, 2, 0, tzinfo=datetime.timezone.utc) + + # Replace the datetime.datetime class with our mock + mock_datetime.datetime = MockDatetime + + # Mock dateparser.parse to return a fixed datetime + with patch("cloudproxy.providers.hetzner.main.dateparser.parse") as mock_parse: + mock_parse.return_value = datetime.datetime(2023, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc) + + # Run the function + result = hetzner_check_alive(mock_instance_config) + + # Should delete only one instance due to rolling deployment min_available rules + mock_delete_proxy.assert_called_once_with(mock_proxy, mock_instance_config) + assert len(result) == 1 @patch("cloudproxy.providers.hetzner.main.config", new_callable=MagicMock) @patch("cloudproxy.providers.hetzner.main.check_alive") diff --git a/tests/test_providers_vultr_main.py b/tests/test_providers_vultr_main.py index af1d4ef..3dc79bf 100644 --- a/tests/test_providers_vultr_main.py +++ b/tests/test_providers_vultr_main.py @@ -41,7 +41,23 @@ def mock_config(self, mock_instance_config): "rolling_deployment": {"enabled": False, "min_available": 3, "batch_size": 2} }.get(x, {}) yield mock_cfg - + + @pytest.fixture + def mock_config_rolling(self, mock_instance_config): + with patch('cloudproxy.providers.vultr.main.config') as mock_cfg: + mock_cfg.__getitem__.side_effect = lambda x: { + "providers": { + "vultr": { + "instances": { + "default": mock_instance_config + } + } + }, + "age_limit": 3600, # 1 hour + "rolling_deployment": {"enabled": True, "min_available": 1, "batch_size": 1} + }.get(x, {}) + yield mock_cfg + @pytest.fixture def mock_instances(self): """Create mock Vultr instances.""" @@ -164,7 +180,40 @@ def test_vultr_check_alive_age_limit(self, mock_parse, mock_list, mock_delete, # Should delete the old instance mock_delete.assert_called_once_with(mock_instance, mock_instance_config) assert len(result) == 0 - + + @patch('cloudproxy.providers.vultr.main.check_alive') + @patch('cloudproxy.providers.vultr.main.delete_proxy') + @patch('cloudproxy.providers.vultr.main.list_instances') + @patch('cloudproxy.providers.vultr.main.dateparser.parse') + def test_vultr_check_alive_age_limit_with_rolling_deployment(self, mock_parse, mock_list, mock_delete, + mock_check_alive, mock_config_rolling, mock_instance_config): + # Setup - instances are older than age limit + mock_instances = [ + VultrInstance({ + "id": "old-instance", + "main_ip": "192.168.1.1", + "status": "active", + "date_created": "2024-01-01T00:00:00Z" + }), + VultrInstance({ + "id": "old-instance-", + "main_ip": "192.168.1.2", + "status": "active", + "date_created": "2024-01-01T00:00:00Z" + }) + ] + mock_list.return_value = mock_instances + # Instance created 2 hours ago, age limit is 1 hour + mock_parse.return_value = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=2) + mock_check_alive.return_value = True + + # Call function + result = vultr_check_alive(mock_instance_config) + + # Should delete only one instance due to rolling deployment min_available rules + mock_delete.assert_called_once_with(mock_instances[0], mock_instance_config) + assert len(result) == 1 + @patch('cloudproxy.providers.vultr.main.check_alive') @patch('cloudproxy.providers.vultr.main.delete_proxy') @patch('cloudproxy.providers.vultr.main.list_instances') diff --git a/tests/test_rolling_deployment.py b/tests/test_rolling_deployment.py index 94e5616..3f3a954 100644 --- a/tests/test_rolling_deployment.py +++ b/tests/test_rolling_deployment.py @@ -174,7 +174,7 @@ def test_update_proxy_health(self): healthy_ips = ["192.168.1.1", "192.168.1.2"] pending_ips = ["192.168.1.3"] - self.manager.update_proxy_health("aws", "default", healthy_ips, pending_ips) + self.manager.update_proxy_health("aws", "default", set(healthy_ips), set(pending_ips)) state = self.manager.get_state("aws", "default") assert state.healthy_proxies == {"192.168.1.1", "192.168.1.2"} @@ -188,7 +188,7 @@ def test_update_proxy_health_cleans_stale_recycling(self): state.pending_recycle.add("192.168.1.3") # No longer exists healthy_ips = ["192.168.1.2"] - self.manager.update_proxy_health("aws", "default", healthy_ips, []) + self.manager.update_proxy_health("aws", "default", set(healthy_ips), set()) assert "192.168.1.1" not in state.recycling assert "192.168.1.2" in state.recycling @@ -246,7 +246,7 @@ def test_complex_rolling_scenario(self): """Test a complex rolling deployment scenario.""" # Initial state: 5 healthy proxies healthy_ips = [f"192.168.1.{i}" for i in range(1, 6)] - self.manager.update_proxy_health("aws", "default", healthy_ips, []) + self.manager.update_proxy_health("aws", "default", set(healthy_ips), set()) # Try to recycle 3 proxies with batch_size=2, min_available=3 results = [] diff --git a/tests/test_rolling_deployment_simple.py b/tests/test_rolling_deployment_simple.py index be9f36c..e0822ea 100644 --- a/tests/test_rolling_deployment_simple.py +++ b/tests/test_rolling_deployment_simple.py @@ -38,8 +38,8 @@ def test_can_recycle_with_sufficient_proxies(self): # Update health status first rolling_manager.update_proxy_health( "test", "default", - healthy_ips=["1.1.1.1", "2.2.2.2", "3.3.3.3", "4.4.4.4"], - pending_ips=[] + healthy_ips=set(["1.1.1.1", "2.2.2.2", "3.3.3.3", "4.4.4.4"]), + pending_ips=set() ) # Test recycling with sufficient proxies @@ -63,8 +63,8 @@ def test_cannot_recycle_below_minimum(self): # Update health status first rolling_manager.update_proxy_health( "test", "default", - healthy_ips=["1.1.1.1", "2.2.2.2"], - pending_ips=[] + healthy_ips=set(["1.1.1.1", "2.2.2.2"]), + pending_ips=set() ) # Test recycling that would go below minimum @@ -88,8 +88,8 @@ def test_batch_size_limit(self): # Update health status first rolling_manager.update_proxy_health( "test", "default", - healthy_ips=["1.1.1.1", "2.2.2.2", "3.3.3.3", "4.4.4.4", "5.5.5.5"], - pending_ips=[] + healthy_ips=set(["1.1.1.1", "2.2.2.2", "3.3.3.3", "4.4.4.4", "5.5.5.5"]), + pending_ips=set() ) # Try to recycle multiple proxies with batch_size=2 @@ -132,8 +132,8 @@ def test_state_tracking(self): # Update proxy health rolling_manager.update_proxy_health( "test", "default", - healthy_ips=["1.1.1.1", "2.2.2.2"], - pending_ips=["3.3.3.3"] + healthy_ips=set(["1.1.1.1", "2.2.2.2"]), + pending_ips=set(["3.3.3.3"]) ) # Check state @@ -163,8 +163,8 @@ def test_min_scaling_adjustment(self): # Update health status first rolling_manager.update_proxy_health( "test", "default", - healthy_ips=["1.1.1.1", "2.2.2.2", "3.3.3.3"], - pending_ips=[] + healthy_ips=set(["1.1.1.1", "2.2.2.2", "3.3.3.3"]), + pending_ips=set() ) # Test with min_scaling that adjusts effective minimum @@ -187,14 +187,14 @@ def test_recycling_status_report(self): # Setup some state rolling_manager.update_proxy_health( "aws", "default", - healthy_ips=["1.1.1.1", "2.2.2.2"], - pending_ips=["3.3.3.3"] + healthy_ips=set(["1.1.1.1", "2.2.2.2"]), + pending_ips=set(["3.3.3.3"]) ) rolling_manager.update_proxy_health( "gcp", "production", - healthy_ips=["4.4.4.4"], - pending_ips=[] + healthy_ips=set(["4.4.4.4"]), + pending_ips=set() ) # Get status @@ -213,8 +213,8 @@ def test_should_create_replacement(self): # Setup state with proxies being recycled rolling_manager.update_proxy_health( "test", "default", - healthy_ips=["1.1.1.1"], - pending_ips=[] + healthy_ips=set(["1.1.1.1"]), + pending_ips=set() ) # Test with min_scaling higher than current count