tests: update tests that presume a single AZ

test_runner/regress/test_sharding.py

@@ -44,6 +44,10 @@ def test_sharding_smoke(
     """
 
     shard_count = 4
+
+    # Single AZ so that a sharded tenant can spread across all pageservers
+    neon_env_builder.num_azs = 1
+
     neon_env_builder.num_pageservers = shard_count
 
     # 1MiB stripes: enable getting some meaningful data distribution without
@@ -523,16 +527,13 @@ def test_sharding_split_smoke(
     # In preferred AZ & other AZ we will end up with one shard per pageserver
     neon_env_builder.num_pageservers = split_shard_count * 2
 
-    # Two AZs
-    def assign_az(ps_cfg):
-        az = f"az-{(ps_cfg['id'] - 1) % 2}"
-        ps_cfg["availability_zone"] = az
-
+    def ps_conf_hook(ps_cfg):
         # We will run more pageservers than tests usually do, so give them tiny page caches
         # in case we're on a test node under memory pressure.
         ps_cfg["page_cache_size"] = 128
 
-    neon_env_builder.pageserver_config_override = assign_az
+    neon_env_builder.num_azs = 2
+    neon_env_builder.pageserver_config_override = ps_conf_hook
 
     # 1MiB stripes: enable getting some meaningful data distribution without
     # writing large quantities of data in this test.  The stripe size is given

test_runner/regress/test_storage_controller.py

@@ -2212,6 +2212,7 @@ def test_graceful_cluster_restart(
     tenants where we fill based on a target shard count.
     """
     neon_env_builder.num_azs = num_azs
+
     neon_env_builder.num_pageservers = 2
     neon_env_builder.control_plane_hooks_api = compute_reconfigure_listener.control_plane_hooks_api
     env = neon_env_builder.init_configs()
@@ -2475,7 +2476,6 @@ def test_storage_controller_node_deletion(
     Test that deleting a node works & properly reschedules everything that was on the node.
     """
     neon_env_builder.num_pageservers = 3
-    neon_env_builder.num_azs = 3
     env = neon_env_builder.init_configs()
     env.start()
 
@@ -3512,12 +3512,7 @@ def eq_safekeeper_records(a: dict[str, Any], b: dict[str, Any]) -> bool:
 
 @run_only_on_default_postgres("this is like a 'unit test' against storcon db")
 def test_shard_preferred_azs(neon_env_builder: NeonEnvBuilder):
-    def assign_az(ps_cfg):
-        az = f"az-{ps_cfg['id'] % 2}"
-        log.info("Assigned AZ {az}")
-        ps_cfg["availability_zone"] = az
-
-    neon_env_builder.pageserver_config_override = assign_az
+    neon_env_builder.num_azs = 2
     neon_env_builder.num_pageservers = 4
     env = neon_env_builder.init_configs()
     env.start()

test_runner/regress/test_timeline_detach_ancestor.py

@@ -1186,6 +1186,8 @@ def test_sharded_tad_interleaved_after_partial_success(neon_env_builder: NeonEnv
     """
 
     shard_count = 2
+    # Single AZ so that a sharded tenant can spread across all pageservers
+    neon_env_builder.num_azs = 1
     neon_env_builder.num_pageservers = shard_count
     env = neon_env_builder.init_start(initial_tenant_shard_count=shard_count)