Create a table in each partition's reporting region, instead of just us-east-1 (#6352)

Author: dreambeyondorange

tests/integration-tests/conftest.py

@@ -40,12 +40,12 @@
 )
 from conftest_networking import unmarshal_az_override
 from conftest_tests_config import apply_cli_dimensions_filtering, parametrize_from_config, remove_disabled_tests
-from conftest_utils import add_filename_markers
+from conftest_utils import add_filename_markers, get_reporting_region
 from constants import SCHEDULERS_SUPPORTING_IMDS_SECURED, NodeType
 from filelock import FileLock
 from framework.credential_providers import aws_credential_provider, register_cli_credentials_for_region
 from framework.fixture_utils import xdist_session_fixture
-from framework.framework_constants import METADATA_DEFAULT_REGION, METADATA_TABLE
+from framework.framework_constants import METADATA_TABLE
 from framework.metadata_table_manager import MetadataTableManager
 from framework.tests_configuration.config_renderer import read_config_file
 from framework.tests_configuration.config_utils import get_all_regions
@@ -258,7 +258,7 @@ def pytest_configure(config):
     _setup_custom_logger(config.getoption("tests_log_file"))
 
 
-def pytest_sessionstart(session):
+def pytest_sessionstart(session: pytest.Session):
     # The number of seconds before a connection to the instance metadata service should time out.
     # When attempting to retrieve credentials on an Amazon EC2 instance that is configured with an IAM role,
     # a connection to the instance metadata service will time out after 1 second by default. If you know you're
@@ -271,11 +271,6 @@ def pytest_sessionstart(session):
     os.environ["AWS_METADATA_SERVICE_NUM_ATTEMPTS"] = "5"
     # Increasing default max attempts retry
     os.environ["AWS_MAX_ATTEMPTS"] = "10"
-    try:
-        # Setup the metadata table in case it doesn't exist
-        MetadataTableManager(METADATA_DEFAULT_REGION, METADATA_TABLE).create_metadata_table()
-    except Exception as exc:
-        logging.info(f"There was a '{type(exc)}' error with '{exc}' when creating the table!")
 
 
 def pytest_fixture_setup(fixturedef: FixtureDef[Any], request: SubRequest) -> Optional[object]:
@@ -309,6 +304,17 @@ def pytest_collection_modifyitems(session: pytest.Session, config: pytest.Config
 
 def pytest_collection_finish(session):
     _log_collected_tests(session)
+    # Create the metadata table after the regions are collected
+    try:
+        region = session.config.getoption("regions") or get_all_regions(session.config.getoption("tests_config"))
+        region = [unmarshal_az_override(az) for az in region]
+        # Use the first element of the list of regions, since there must be at least one
+        reporting_region = get_reporting_region(region[0])
+        logging.info(f"Metadata reporting region {reporting_region}")
+        # Setup the metadata table in case it doesn't exist
+        MetadataTableManager(reporting_region, METADATA_TABLE).create_metadata_table()
+    except Exception as exc:
+        logging.info(f"There was a '{type(exc)}' error with '{exc}' when creating the table!")
 
 
 def _log_collected_tests(session):

tests/integration-tests/conftest_utils.py

@@ -23,11 +23,18 @@
 import yaml
 from conftest_markers import DIMENSIONS_MARKER_ARGS
 from filelock import FileLock
-from framework.framework_constants import METADATA_DEFAULT_REGION, METADATA_TABLE
+from framework.framework_constants import METADATA_TABLE
 from framework.metadata_table_manager import MetadataTableManager, PhaseMetadata, TestMetadata
 from framework.metrics_publisher import Metric, MetricsPublisher
 from time_utils import microseconds
-from utils import dict_add_nested_key, dict_has_nested_key
+from utils import (
+    DEFAULT_PARTITION,
+    DEFAULT_REPORTING_REGION,
+    PARTITION_MAP,
+    REPORTING_REGION_MAP,
+    dict_add_nested_key,
+    dict_has_nested_key,
+)
 
 
 def add_properties_to_report(item: pytest.Item):
@@ -169,8 +176,9 @@ def create_phase_metrics(item: pytest.Item, rep: pytest.TestReport, dimensions:
 
 def publish_test_metadata(item: pytest.Item, rep: pytest.TestReport):
     """Publish test metadata to the metadata table."""
-    metadata_table_mgr = MetadataTableManager(METADATA_DEFAULT_REGION, METADATA_TABLE)
-    logging.info(f"Publishing test metadata: item {item} rep {rep}")
+    reporting_region = get_reporting_region(get_user_prop(item, "region"))
+    metadata_table_mgr = MetadataTableManager(reporting_region, METADATA_TABLE)
+    logging.info(f"Publishing test metadata: item {item} rep {rep} to {METADATA_TABLE} in {reporting_region}")
     test_metadata = None
     if rep.when == "setup":
         # Initialize the test data
@@ -218,3 +226,15 @@ def publish_test_metadata(item: pytest.Item, rep: pytest.TestReport):
         item.user_properties.append(("metadata", jsonpickle.encode(test_metadata)))
         logging.info(f"Added the metadata during the {rep.when} phase: {get_user_prop(item, 'metadata')}")
     metadata_table_mgr.publish_metadata([test_metadata])
+
+
+def get_reporting_region(region: str):
+    """Get partition for the given region. If region is None, consider the region set in the environment."""
+    curr_partition = next(
+        (partition for region_prefix, partition in PARTITION_MAP.items() if region.startswith(region_prefix)),
+        DEFAULT_PARTITION,
+    )
+    return next(
+        (region for partition, region in REPORTING_REGION_MAP.items() if partition == curr_partition),
+        DEFAULT_REPORTING_REGION,
+    )

tests/integration-tests/utils.py

@@ -36,6 +36,14 @@
     "us-iso-": "aws-iso",
     "us-isob": "aws-iso-b",
 }
+DEFAULT_REPORTING_REGION = "us-east-1"
+# This map is based on where jenkins is deployed
+REPORTING_REGION_MAP = {
+    "aws-cn": "cn-north-1",
+    "aws-us-gov": "us-gov-west-1",
+    "aws-iso": "us-iso-east-1",
+    "aws-iso-b": "us-isob-east-1",
+}
 
 
 def _format_stack_error(message, stack_events=None, cluster_details=None) -> str: