From 3599b250542a39e455ccf5b6528ae902c9343ad4 Mon Sep 17 00:00:00 2001
From: shengweiw <shengweiw@netflix.com>
Date: Fri, 15 Sep 2023 11:32:23 -0700
Subject: [PATCH] Determine postgres region topology by using num_regions
 instead of extra arguments

---
 .../models/org/netflix/postgres.py                    |  2 +-
 tests/netflix/test_postgres.py                        | 11 ++++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/service_capacity_modeling/models/org/netflix/postgres.py b/service_capacity_modeling/models/org/netflix/postgres.py
index db4e3b4..153ec99 100644
--- a/service_capacity_modeling/models/org/netflix/postgres.py
+++ b/service_capacity_modeling/models/org/netflix/postgres.py
@@ -32,7 +32,7 @@ def capacity_plan(
         if desires.service_tier == 0:
             return None
 
-        require_multi_region: bool = extra_model_arguments.get("require_multi_region", False)
+        require_multi_region: bool = context.num_regions > 1
         plan = None
         if Platform.aurora_postgres in instance.platforms and not require_multi_region:
             plan = nflx_aurora_capacity_model.capacity_plan(
diff --git a/tests/netflix/test_postgres.py b/tests/netflix/test_postgres.py
index 28553e5..1a85269 100644
--- a/tests/netflix/test_postgres.py
+++ b/tests/netflix/test_postgres.py
@@ -67,6 +67,7 @@ def test_tier_0_not_supported():
         model_name="org.netflix.postgres",
         region="us-east-1",
         desires=tier_0,
+        num_regions=1,
     )
     # Aurora can't support tier 0 service
     assert len(cap_plan) == 0
@@ -77,6 +78,7 @@ def test_small_footprint():
         model_name="org.netflix.postgres",
         region="us-east-1",
         desires=small_footprint,
+        num_regions=1,
     )
     assert cap_plan[0].candidate_clusters.regional[0].instance.name == "db.r5.large"
 
@@ -95,7 +97,7 @@ def test_small_footprint_multi_region():
         model_name="org.netflix.postgres",
         region="us-east-1",
         desires=small_footprint,
-        extra_model_arguments={"require_multi_region": True}
+        num_regions=3,
     )
     assert cap_plan[0].candidate_clusters.zonal[0].instance.name == "m5d.xlarge"
 
@@ -105,16 +107,13 @@ def test_small_footprint_multi_region():
         < 4000
     )
 
-    planner.plan(model_name="org.netflix.postgres", region="us-east-1", desires=small_footprint,
-                 extra_model_arguments={"require_multi_region": False})
-
 
 def test_small_footprint_plan_uncertain():
     cap_plan = planner.plan(
         model_name="org.netflix.postgres",
         region="us-east-1",
         desires=small_footprint,
-        extra_model_arguments={"require_multi_region": False},
+        num_regions=1,
         simulations=256
     )
     plan_a = cap_plan.least_regret[0]
@@ -133,6 +132,7 @@ def test_large_footprint():
         model_name="org.netflix.postgres",
         region="us-east-1",
         desires=large_footprint,
+        num_regions=1,
     )
     # Aurora cannot handle the scale, so pick crdb
     assert cap_plan[0].candidate_clusters.zonal[0].instance.name == "i3.xlarge"
@@ -150,5 +150,6 @@ def test_tier_3():
         model_name="org.netflix.postgres",
         region="us-east-1",
         desires=tier_3,
+        num_regions=1,
     )
     assert cap_plan[0].candidate_clusters.regional[0].instance.name == "db.r5.2xlarge"