Using params in default_desires to determine repl factor

Author: rayiniv-nflx

service_capacity_modeling/models/org/netflix/kafka.py

@@ -508,10 +508,11 @@ def default_desires(user_desires, extra_model_arguments: Dict[str, Any]):
         retention_secs = iso_to_seconds(retention)
 
         # write throughput * retention * replication factor = usage
+        replication_factor = NflxKafkaCapacityModel.HA_DEFAULT_REPLICATION_FACTOR
+        if extra_model_arguments.get("cluster_type", None) == ClusterType.strong:
+            replication_factor = NflxKafkaCapacityModel.SC_DEFAULT_REPLICATION_FACTOR
         state_gib = (
-            write_bytes.mid
-            * retention_secs
-            * NflxKafkaCapacityModel.HA_DEFAULT_REPLICATION_FACTOR
+            write_bytes.mid * retention_secs * replication_factor
         ) / GIB_IN_BYTES
 
         return CapacityDesires(

tests/netflix/test_kafka.py

@@ -1,15 +1,13 @@
 from service_capacity_modeling.capacity_planner import planner
-from service_capacity_modeling.interface import (
-    CapacityDesires,
-    AccessPattern,
-    FixedInterval,
-    DataShape,
-    Drive,
-    DriveType,
-)
+from service_capacity_modeling.interface import AccessPattern
+from service_capacity_modeling.interface import CapacityDesires
 from service_capacity_modeling.interface import certain_float
 from service_capacity_modeling.interface import CurrentClusters
 from service_capacity_modeling.interface import CurrentZoneClusterCapacity
+from service_capacity_modeling.interface import DataShape
+from service_capacity_modeling.interface import Drive
+from service_capacity_modeling.interface import DriveType
+from service_capacity_modeling.interface import FixedInterval
 from service_capacity_modeling.interface import Interval
 from service_capacity_modeling.interface import QueryPattern
 from service_capacity_modeling.models.common import normalize_cores
@@ -72,7 +70,7 @@ def test_kafka_large_scale():
         region="us-east-1",
         desires=desires,
         extra_model_arguments={
-            "cluster_type": "strong",
+            "cluster_type": ClusterType.strong,
             "retention": "PT4H",
         },
     )