diff --git a/codecarbon/external/geography.py b/codecarbon/external/geography.py index 075824959..6b211a61a 100644 --- a/codecarbon/external/geography.py +++ b/codecarbon/external/geography.py @@ -28,7 +28,7 @@ def extract_gcp_region(zone: str) -> str: """ projects/705208488469/zones/us-central1-a -> us-central1 """ - google_region_regex = r"[a-z]+-[a-z]+[0-9]" + google_region_regex = r"[a-z]+-[a-z]+[0-9]+" return re.search(google_region_regex, zone).group(0) extract_region_for_provider: Dict[str, Callable] = { diff --git a/tests/test_geography.py b/tests/test_geography.py index 8f95f7f43..b0a61cacb 100644 --- a/tests/test_geography.py +++ b/tests/test_geography.py @@ -1,6 +1,7 @@ import unittest from unittest import mock +import pytest import responses from codecarbon.external.geography import CloudMetadata, GeoMetadata @@ -15,6 +16,20 @@ ) +@pytest.mark.parametrize( + "region", ["europe-west12", "europe-west10", "us-central1", "asia-northeast3"] +) +@pytest.mark.parametrize("prefix", ["", "projects/123/zones/"]) +def test_cloud_metadata_gcp_region_number(region, prefix): + metadata = {"provider": "gcp", "metadata": {"zone": f"{prefix}{region}-a"}} + with mock.patch( + "codecarbon.external.geography.get_env_cloud_details", return_value=metadata + ): + cloud = CloudMetadata.from_utils() + assert cloud.provider == "gcp" + assert cloud.region == region + + class TestCloudMetadata(unittest.TestCase): @mock.patch( "codecarbon.external.geography.get_env_cloud_details",