service-capacity-modeling 0.3.88__tar.gz → 0.3.90__tar.gz

{service_capacity_modeling-0.3.88 → service_capacity_modeling-0.3.90}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: service-capacity-modeling
-Version: 0.3.88
+Version: 0.3.90
 Summary: Contains utilities for modeling capacity for pluggable workloads
 Author: Joseph Lynch
 Author-email: josephl@netflix.com

{service_capacity_modeling-0.3.88 → service_capacity_modeling-0.3.90}/service_capacity_modeling/enum_utils.py

@@ -5,6 +5,7 @@ runtime-accessible docstrings.
 
 import ast
 import inspect
+import sys
 from enum import Enum
 from functools import partial
 from operator import is_
@@ -16,6 +17,47 @@ from pydantic.json_schema import JsonSchemaValue
 from pydantic_core import CoreSchema
 
 
+__all__ = ["StrEnum", "enum_docstrings"]
+
+# StrEnum backport for Python 3.10 compatibility
+# On Python 3.11+, use the stdlib version
+if sys.version_info >= (3, 11):
+    from enum import StrEnum as StrEnum  # pylint: disable=useless-import-alias
+else:
+
+    class StrEnum(str, Enum):
+        """Backport of Python 3.11 StrEnum.
+
+        Provides consistent string behavior across all Python versions:
+        - f"{x}" returns the value (not "Foo.BAR")
+        - str(x) returns the value (not "Foo.BAR")
+        - x == "value" returns True (string comparison works)
+
+        This addresses PEP 663 which changed str(Enum) behavior in Python 3.11,
+        making (str, Enum) return "Foo.BAR" in f-strings instead of the value.
+        """
+
+        def __new__(cls, value: str, *args: Any, **kwargs: Any) -> "StrEnum":
+            if not isinstance(value, str):
+                raise TypeError(f"{value!r} is not a string")
+            member = str.__new__(cls, value)
+            member._value_ = value
+            return member
+
+        def __str__(self) -> str:
+            return str(self.value)
+
+        def __format__(self, format_spec: str) -> str:
+            # Ensures f-strings return value, not "Foo.BAR"
+            return str(self.value).__format__(format_spec)
+
+        @staticmethod
+        def _generate_next_value_(
+            name: str, start: int, count: int, last_values: list[str]
+        ) -> str:
+            return name.lower()
+
+
 E = TypeVar("E", bound=Enum)
 
 

{service_capacity_modeling-0.3.88 → service_capacity_modeling-0.3.90}/service_capacity_modeling/interface.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 import re
 import sys
 from decimal import Decimal
-from enum import Enum
 from fractions import Fraction
 from functools import lru_cache
 from typing import Any
@@ -22,6 +21,7 @@ from pydantic import ConfigDict
 from pydantic import Field
 
 from service_capacity_modeling.enum_utils import enum_docstrings
+from service_capacity_modeling.enum_utils import StrEnum
 
 GIB_IN_BYTES = 1024 * 1024 * 1024
 MIB_IN_BYTES = 1024 * 1024
@@ -46,16 +46,13 @@ class ExcludeUnsetModel(BaseModel):
 
 
 @enum_docstrings
-class IntervalModel(str, Enum):
+class IntervalModel(StrEnum):
     """Statistical distribution models for approximating intervals
 
     When we have uncertainty intervals (low, mid, high), we need to choose
     a probability distribution to model that uncertainty for simulation purposes.
     """
 
-    def __str__(self) -> str:
-        return str(self.value)
-
     def __repr__(self) -> str:
         return f"D({self.value})"
 
@@ -204,7 +201,7 @@ def normalized_aws_size(name: str) -> Fraction:
 ###############################################################################
 
 
-class Lifecycle(str, Enum):
+class Lifecycle(StrEnum):
     """Represents the lifecycle of hardware from initial preview
     to end-of-life.
 
@@ -223,7 +220,7 @@ class Lifecycle(str, Enum):
 
 
 @enum_docstrings
-class DriveType(str, Enum):
+class DriveType(StrEnum):
     """Represents the type and attachment model of storage drives
 
     Drives can be either local (ephemeral, instance-attached) or network-attached
@@ -366,7 +363,7 @@ class Drive(ExcludeUnsetModel):
 
 
 @enum_docstrings
-class Platform(str, Enum):
+class Platform(StrEnum):
     """Represents the CPU architecture or managed service platform
 
     Hardware can run on different CPU architectures (x86_64, ARM) or be a fully
@@ -617,7 +614,7 @@ class Pricing(ExcludeUnsetModel):
 
 
 @enum_docstrings
-class AccessPattern(str, Enum):
+class AccessPattern(StrEnum):
     """The access pattern determines capacity planning priorities: latency-sensitive
     services target low P99 latency, while throughput-oriented services optimize
     for maximum requests per second.
@@ -633,7 +630,7 @@ class AccessPattern(str, Enum):
 
 
 @enum_docstrings
-class AccessConsistency(str, Enum):
+class AccessConsistency(StrEnum):
     """
     Generally speaking consistency is expensive, so models need to know what
     kind of consistency will be required in order to estimate CPU usage
@@ -851,7 +848,7 @@ class CurrentClusters(ExcludeUnsetModel):
 
 
 @enum_docstrings
-class BufferComponent(str, Enum):
+class BufferComponent(StrEnum):
     """Represents well known buffer components such as compute and storage
 
     Note that while these are common and defined here for models to share,
@@ -887,7 +884,7 @@ class BufferComponent(str, Enum):
 
 
 @enum_docstrings
-class BufferIntent(str, Enum):
+class BufferIntent(StrEnum):
     """Defines the intent of buffer directives for capacity planning"""
 
     desired = "desired"

{service_capacity_modeling-0.3.88 → service_capacity_modeling-0.3.90}/service_capacity_modeling/models/org/netflix/cassandra.py

@@ -638,8 +638,14 @@ def _target_rf(desires: CapacityDesires, user_copies: Optional[int]) -> int:
 
 
 class NflxCassandraArguments(BaseModel):
-    copies_per_region: int = Field(
-        default=3,
+    """Configuration arguments for the Netflix Cassandra capacity model.
+
+    This model centralizes all tunable parameters with their defaults.
+    Use `from_extra_model_arguments()` to parse a dict into a validated instance.
+    """
+
+    copies_per_region: Optional[int] = Field(
+        default=None,
         description="How many copies of the data will exist e.g. RF=3. If unsupplied"
         " this will be deduced from durability and consistency desires",
     )
@@ -663,9 +669,13 @@ class NflxCassandraArguments(BaseModel):
         default=192,
         description="What is the maximum size of a cluster in this region",
     )
-    max_local_disk_gib: int = Field(
-        default=5120,
-        description="The maximum amount of data we store per machine",
+    max_local_data_per_node_gib: int = Field(
+        default=1280,
+        description="Maximum data per node for local disk instances (GiB)",
+    )
+    max_attached_data_per_node_gib: int = Field(
+        default=2048,
+        description="Maximum data per node for attached disk instances (GiB)",
     )
     max_write_buffer_percent: float = Field(
         default=0.25,
@@ -680,6 +690,26 @@ class NflxCassandraArguments(BaseModel):
         "automatically adjust to 0.2",
     )
 
+    @classmethod
+    def from_extra_model_arguments(
+        cls, extra_model_arguments: Dict[str, Any]
+    ) -> "NflxCassandraArguments":
+        """Parse extra_model_arguments dict into a validated NflxCassandraArguments.
+
+        This centralizes default values - any field not in extra_model_arguments
+        will use the default defined in this model.
+
+        Handles legacy field name mappings:
+        - max_local_disk_gib -> max_local_data_per_node_gib (if not explicitly set)
+        """
+        # Handle legacy field name: max_local_disk_gib -> max_local_data_per_node_gib
+        args = dict(extra_model_arguments)
+        if "max_local_data_per_node_gib" not in args and "max_local_disk_gib" in args:
+            args["max_local_data_per_node_gib"] = args["max_local_disk_gib"]
+
+        # Pydantic will use defaults for any missing fields
+        return cls.model_validate(args)
+
 
 class NflxCassandraCapacityModel(CapacityModel):
     def __init__(self) -> None:
@@ -722,40 +752,22 @@ class NflxCassandraCapacityModel(CapacityModel):
         desires: CapacityDesires,
         extra_model_arguments: Dict[str, Any],
     ) -> Optional[CapacityPlan]:
-        # TODO: Standardize these extra model argument defaults in a single
-        # place. Many of them are defined here and as default values in the
-        # downstream method but only these ones are used which is confusing for
-        # readability
-
-        # Use durabiliy and consistency to compute RF.
-        copies_per_region = _target_rf(
-            desires, extra_model_arguments.get("copies_per_region", None)
-        )
-        require_local_disks: bool = extra_model_arguments.get(
-            "require_local_disks", False
-        )
-        require_attached_disks: bool = extra_model_arguments.get(
-            "require_attached_disks", False
-        )
+        # Parse extra_model_arguments into a validated model with centralized defaults
+        args = NflxCassandraArguments.from_extra_model_arguments(extra_model_arguments)
+
+        # Use durability and consistency to compute RF if not explicitly set
+        copies_per_region = _target_rf(desires, args.copies_per_region)
+
+        # Validate required_cluster_size for critical tiers
         required_cluster_size: Optional[int] = (
             NflxCassandraCapacityModel.get_required_cluster_size(
                 desires.service_tier, extra_model_arguments
             )
         )
 
-        max_rps_to_disk: int = extra_model_arguments.get("max_rps_to_disk", 500)
-        max_regional_size: int = extra_model_arguments.get("max_regional_size", 192)
-        max_local_data_per_node_gib: int = extra_model_arguments.get(
-            "max_local_data_per_node_gib",
-            extra_model_arguments.get("max_local_disk_gib", 1280),
-        )
-
-        max_write_buffer_percent: float = min(
-            0.5, extra_model_arguments.get("max_write_buffer_percent", 0.25)
-        )
-        max_table_buffer_percent: float = min(
-            0.5, extra_model_arguments.get("max_table_buffer_percent", 0.11)
-        )
+        # Apply caps to buffer percentages
+        max_write_buffer_percent = min(0.5, args.max_write_buffer_percent)
+        max_table_buffer_percent = min(0.5, args.max_table_buffer_percent)
 
         # Adjust heap defaults for high write clusters
         if (
@@ -772,12 +784,12 @@ class NflxCassandraCapacityModel(CapacityModel):
             desires=desires,
             zones_per_region=context.zones_in_region,
             copies_per_region=copies_per_region,
-            require_local_disks=require_local_disks,
-            require_attached_disks=require_attached_disks,
+            require_local_disks=args.require_local_disks,
+            require_attached_disks=args.require_attached_disks,
             required_cluster_size=required_cluster_size,
-            max_rps_to_disk=max_rps_to_disk,
-            max_regional_size=max_regional_size,
-            max_local_data_per_node_gib=max_local_data_per_node_gib,
+            max_rps_to_disk=args.max_rps_to_disk,
+            max_regional_size=args.max_regional_size,
+            max_local_data_per_node_gib=args.max_local_data_per_node_gib,
             max_write_buffer_percent=max_write_buffer_percent,
             max_table_buffer_percent=max_table_buffer_percent,
         )

{service_capacity_modeling-0.3.88 → service_capacity_modeling-0.3.90}/service_capacity_modeling/models/org/netflix/counter.py

@@ -1,5 +1,4 @@
 from datetime import timedelta
-from enum import Enum
 from typing import Any
 from typing import Callable
 from typing import Dict
@@ -10,6 +9,7 @@ from pydantic import Field
 
 from .stateless_java import nflx_java_app_capacity_model
 from .stateless_java import NflxJavaAppArguments
+from service_capacity_modeling.enum_utils import StrEnum
 from service_capacity_modeling.interface import AccessConsistency
 from service_capacity_modeling.interface import AccessPattern
 from service_capacity_modeling.interface import CapacityDesires
@@ -26,13 +26,13 @@ from service_capacity_modeling.interface import RegionContext
 from service_capacity_modeling.models import CapacityModel
 
 
-class NflxCounterCardinality(Enum):
+class NflxCounterCardinality(StrEnum):
     low = "low"
     medium = "medium"
     high = "high"
 
 
-class NflxCounterMode(Enum):
+class NflxCounterMode(StrEnum):
     best_effort = "best-effort"
     eventual = "eventual"
     exact = "exact"
@@ -95,7 +95,7 @@ class NflxCounterCapacityModel(CapacityModel):
     ) -> Tuple[Tuple[str, Callable[[CapacityDesires], CapacityDesires]], ...]:
         stores = []
 
-        if extra_model_arguments["counter.mode"] == NflxCounterMode.best_effort.value:
+        if extra_model_arguments["counter.mode"] == NflxCounterMode.best_effort:
             stores.append(("org.netflix.evcache", lambda x: x))
         else:
             # Shared evcache cluster is used for eventual and exact counters
@@ -114,9 +114,9 @@ class NflxCounterCapacityModel(CapacityModel):
                 # high cardinality : rollups happen once every 10 seconds
                 # TODO: Account for read amplification from time slice configs
                 #       for better model accuracy
-                if counter_cardinality == NflxCounterCardinality.low.value:
+                if counter_cardinality == NflxCounterCardinality.low:
                     rollups_per_second = counter_deltas_per_second.scale(0.0167)
-                elif counter_cardinality == NflxCounterCardinality.medium.value:
+                elif counter_cardinality == NflxCounterCardinality.medium:
                     rollups_per_second = counter_deltas_per_second.scale(0.0333)
                 else:
                     rollups_per_second = counter_deltas_per_second.scale(0.1)

{service_capacity_modeling-0.3.88 → service_capacity_modeling-0.3.90}/service_capacity_modeling/models/org/netflix/evcache.py

@@ -1,6 +1,5 @@
 import logging
 import math
-from enum import Enum
 from typing import Any
 from typing import Dict
 from typing import Optional
@@ -9,6 +8,7 @@ from typing import Tuple
 from pydantic import BaseModel
 from pydantic import Field
 
+from service_capacity_modeling.enum_utils import StrEnum
 from service_capacity_modeling.interface import AccessConsistency
 from service_capacity_modeling.interface import AccessPattern
 from service_capacity_modeling.interface import Buffer
@@ -45,7 +45,7 @@ from service_capacity_modeling.stats import dist_for_interval
 logger = logging.getLogger(__name__)
 
 
-class Replication(str, Enum):
+class Replication(StrEnum):
     none = "none"
     sets = "sets"
     evicts = "evicts"

{service_capacity_modeling-0.3.88 → service_capacity_modeling-0.3.90}/service_capacity_modeling/models/org/netflix/kafka.py

@@ -1,6 +1,5 @@
 import logging
 import math
-from enum import Enum
 from typing import Any
 from typing import Dict
 from typing import Optional
@@ -9,6 +8,7 @@ from typing import Tuple
 from pydantic import BaseModel
 from pydantic import Field
 
+from service_capacity_modeling.enum_utils import StrEnum
 from service_capacity_modeling.interface import AccessConsistency
 from service_capacity_modeling.interface import AccessPattern
 from service_capacity_modeling.interface import Buffer
@@ -45,7 +45,7 @@ from service_capacity_modeling.models.org.netflix.iso_date_math import iso_to_se
 logger = logging.getLogger(__name__)
 
 
-class ClusterType(str, Enum):
+class ClusterType(StrEnum):
     strong = "strong"
     ha = "high-availability"
 

{service_capacity_modeling-0.3.88 → service_capacity_modeling-0.3.90}/service_capacity_modeling.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: service-capacity-modeling
-Version: 0.3.88
+Version: 0.3.90
 Summary: Contains utilities for modeling capacity for pluggable workloads
 Author: Joseph Lynch
 Author-email: josephl@netflix.com

service_capacity_modeling-0.3.90/tests/test_enum_utils.py

@@ -0,0 +1,316 @@
+import pytest
+from pydantic import BaseModel
+from pydantic import ValidationError
+
+from service_capacity_modeling.interface import AccessConsistency
+from service_capacity_modeling.interface import AccessPattern
+from service_capacity_modeling.interface import BufferComponent
+from service_capacity_modeling.interface import BufferIntent
+from service_capacity_modeling.interface import DriveType
+from service_capacity_modeling.interface import IntervalModel
+from service_capacity_modeling.interface import Lifecycle
+from service_capacity_modeling.interface import Platform
+from service_capacity_modeling.models.org.netflix.counter import NflxCounterCardinality
+from service_capacity_modeling.models.org.netflix.counter import NflxCounterMode
+from service_capacity_modeling.models.org.netflix.evcache import Replication
+from service_capacity_modeling.models.org.netflix.kafka import ClusterType
+
+# List of all enums that should have per-member docstrings
+DOCUMENTED_ENUMS = [
+    IntervalModel,
+    DriveType,
+    Platform,
+    AccessPattern,
+    AccessConsistency,
+    BufferComponent,
+    BufferIntent,
+]
+
+
+@pytest.mark.parametrize("enum_class", DOCUMENTED_ENUMS)
+def test_enums_have_docstrings(enum_class):
+    """Test that all interface.py enums have comprehensive per-member
+    docstrings
+
+    This test ensures that all enum members have their own
+    runtime-accessible docstrings, which makes them discoverable via
+    help(), IDE tooltips, and the __doc__ attribute.
+
+    The enums use the @enum_docstrings decorator which parses source code
+    to attach docstrings that appear below each member (following PEP 257
+    attribute docstring conventions).
+
+    See: https://stackoverflow.com/questions/19330460/
+    how-do-i-put-docstrings-on-enums
+    """
+    enum_name = enum_class.__name__
+
+    # Check class has a docstring
+    assert enum_class.__doc__ is not None, f"{enum_name} must have a class docstring"
+    assert len(enum_class.__doc__) > 0, f"{enum_name} class docstring must not be empty"
+
+    # Check each member has its own unique docstring
+    for member in enum_class:
+        assert member.__doc__ is not None, (
+            f"{enum_name}.{member.name} must have a docstring. "
+            f"Add a docstring after the member definition:\n"
+            f'    {member.name} = "{member.value}"\n'
+            f'    """Your documentation here"""'
+        )
+        assert len(member.__doc__.strip()) > 0, (
+            f"{enum_name}.{member.name} docstring must not be empty"
+        )
+        # Verify it's not just the class docstring (should be member-specific)
+        assert member.__doc__ != enum_class.__doc__, (
+            f"{enum_name}.{member.name} should have its own docstring, "
+            f"not inherit the class docstring. "
+            f"Did the @enum_docstrings decorator work?"
+        )
+
+    # Verify different members have different docstrings
+    members = list(enum_class)
+    if len(members) >= 2:
+        assert members[0].__doc__ != members[1].__doc__, (
+            f"{enum_name}: Different enum members should have different docstrings"
+        )
+
+
+@pytest.mark.parametrize("enum_class", DOCUMENTED_ENUMS)
+def test_enums_json_schema_includes_member_docstrings(enum_class):
+    """Test that enum member docstrings appear in Pydantic JSON schemas
+
+    The @enum_docstrings decorator adds __get_pydantic_json_schema__ to
+    generate oneOf schemas with per-member descriptions. This ensures
+    enum documentation is available in API schemas, OpenAPI specs, etc.
+    """
+    enum_name = enum_class.__name__
+
+    # Create a test model using this enum
+    TestModel = type(
+        "TestModel",
+        (BaseModel,),
+        {"__annotations__": {"field": enum_class}},
+    )
+
+    # Get the JSON schema
+    schema = TestModel.model_json_schema()
+
+    # Check the enum definition exists in $defs
+    assert "$defs" in schema, f"{enum_name}: Schema missing $defs"
+    assert enum_name in schema["$defs"], f"{enum_name}: Enum not in $defs"
+
+    enum_schema = schema["$defs"][enum_name]
+
+    # Check oneOf exists with member descriptions
+    assert "oneOf" in enum_schema, (
+        f"{enum_name}: JSON schema missing oneOf for member descriptions"
+    )
+
+    one_of = enum_schema["oneOf"]
+    assert len(one_of) == len(enum_class), (
+        f"{enum_name}: oneOf should have {len(enum_class)} entries"
+    )
+
+    # Verify each member has proper schema entry
+    for member in enum_class:
+        matching_entries = [
+            entry for entry in one_of if entry.get("const") == member.value
+        ]
+
+        assert len(matching_entries) == 1, (
+            f"{enum_name}.{member.name}: Should have exactly one oneOf entry"
+        )
+
+        entry = matching_entries[0]
+
+        # Check required fields
+        assert "const" in entry, f"{enum_name}.{member.name}: Missing 'const'"
+        assert "title" in entry, f"{enum_name}.{member.name}: Missing 'title'"
+        assert "description" in entry, (
+            f"{enum_name}.{member.name}: Missing 'description'"
+        )
+
+        # Verify description matches member docstring
+        assert entry["description"] == member.__doc__, (
+            f"{enum_name}.{member.name}: Schema description doesn't match "
+            f"member.__doc__"
+        )
+
+        # Verify description is not empty and not the class docstring
+        assert len(entry["description"].strip()) > 0, (
+            f"{enum_name}.{member.name}: Description should not be empty"
+        )
+        assert entry["description"] != enum_class.__doc__, (
+            f"{enum_name}.{member.name}: Description should be member-specific, "
+            f"not the class docstring"
+        )
+
+
+###############################################################################
+#                    StrEnum Behavior Tests (PEP 663)                         #
+###############################################################################
+#
+# These tests validate that StrEnum provides consistent string behavior across
+# all Python versions (3.10, 3.11, 3.12).
+#
+# Background: PEP 663 changed (str, Enum) behavior in Python 3.11:
+# - Python 3.10: f"{x}" returns value, str(x) returns "Foo.BAR"
+# - Python 3.11: f"{x}" returns "Foo.BAR", str(x) returns "Foo.BAR"
+#
+# StrEnum provides consistent behavior where f"{x}", str(x), and x.value ALL
+# return the value string, making enum usage predictable across Python versions.
+#
+# See: https://peps.python.org/pep-0663/
+
+
+# All StrEnum classes in the codebase
+STRENUM_CLASSES = [
+    IntervalModel,
+    Lifecycle,
+    DriveType,
+    Platform,
+    AccessPattern,
+    AccessConsistency,
+    BufferComponent,
+    BufferIntent,
+    NflxCounterCardinality,
+    NflxCounterMode,
+    Replication,
+    ClusterType,
+]
+
+
+@pytest.mark.parametrize("enum_class", STRENUM_CLASSES)
+def test_strenum_inherits_from_str(enum_class):
+    """Test that all enum classes inherit from str (StrEnum behavior).
+
+    This ensures the enum can be used directly in string contexts.
+    """
+    for member in enum_class:
+        assert isinstance(member, str), (
+            f"{enum_class.__name__}.{member.name} should be a str instance"
+        )
+
+
+@pytest.mark.parametrize("enum_class", STRENUM_CLASSES)
+def test_strenum_fstring_returns_value(enum_class):
+    """Test that f-strings return the enum value, not 'EnumName.member'.
+
+    This is the key behavior that changed in Python 3.11 (PEP 663).
+    Without StrEnum, f"{x}" returns "Foo.BAR" instead of "bar".
+    """
+    for member in enum_class:
+        fstring_result = f"{member}"
+        assert fstring_result == member.value, (
+            f'f"{{{{member}}}}" for {enum_class.__name__}.{member.name} returned '
+            f'"{fstring_result}", expected "{member.value}"'
+        )
+
+
+@pytest.mark.parametrize("enum_class", STRENUM_CLASSES)
+def test_strenum_str_returns_value(enum_class):
+    """Test that str(x) returns the enum value, not 'EnumName.member'.
+
+    With StrEnum, str(x) should return the value for consistent behavior.
+    """
+    for member in enum_class:
+        str_result = str(member)
+        assert str_result == member.value, (
+            f"str({enum_class.__name__}.{member.name}) returned "
+            f'"{str_result}", expected "{member.value}"'
+        )
+
+
+@pytest.mark.parametrize("enum_class", STRENUM_CLASSES)
+def test_strenum_format_returns_value(enum_class):
+    """Test that format()/format spec returns the enum value.
+
+    This uses "{}".format() which should behave like f-strings.
+    """
+    for member in enum_class:
+        format_result = "{}".format(member)  # pylint: disable=consider-using-f-string
+        assert format_result == member.value, (
+            f'"{{}}".format({enum_class.__name__}.{member.name}) returned '
+            f'"{format_result}", expected "{member.value}"'
+        )
+
+
+@pytest.mark.parametrize("enum_class", STRENUM_CLASSES)
+def test_strenum_equals_string(enum_class):
+    """Test that enum members compare equal to their string values.
+
+    This is critical for usage with Pydantic model_dump() and dict comparisons.
+    """
+    for member in enum_class:
+        assert member == member.value, (
+            f"{enum_class.__name__}.{member.name} == {member.value!r} should be True"
+        )
+
+
+@pytest.mark.parametrize("enum_class", STRENUM_CLASSES)
+def test_strenum_works_as_dict_key_with_string_lookup(enum_class):
+    """Test that enum members work as dict keys with string lookup.
+
+    When model_dump() returns enum objects, we should be able to use strings
+    to look up values in dicts keyed by those enums.
+    """
+    for member in enum_class:
+        test_dict = {member: "test_value"}
+        # String lookup should work because member IS a string
+        assert test_dict.get(member.value) == "test_value", (
+            f"Dict keyed by {enum_class.__name__}.{member.name} should be "
+            f"accessible via string {member.value!r}"
+        )
+
+
+def test_strenum_pydantic_validation_accepts_valid_strings():
+    """Test that Pydantic accepts valid enum string values."""
+
+    class TestModel(BaseModel):
+        pattern: AccessPattern
+        consistency: AccessConsistency
+
+    # Should accept string values
+    model = TestModel(pattern="latency", consistency="eventual")
+    assert model.pattern == AccessPattern.latency
+    assert model.consistency == AccessConsistency.eventual
+
+
+def test_strenum_pydantic_validation_rejects_invalid_strings():
+    """Test that Pydantic rejects invalid enum string values.
+
+    This ensures strict validation - arbitrary strings are NOT accepted.
+    """
+
+    class TestModel(BaseModel):
+        pattern: AccessPattern
+
+    with pytest.raises(ValidationError):
+        TestModel(pattern="invalid_pattern_value")
+
+
+def test_strenum_pydantic_model_dump_preserves_enum():
+    """Test that model_dump() returns enum objects that behave as strings.
+
+    By default, Pydantic returns enum objects (not raw strings) from model_dump().
+    With StrEnum, these objects ARE strings, so comparisons work correctly.
+    """
+
+    class TestModel(BaseModel):
+        pattern: AccessPattern
+        drive: DriveType
+
+    model = TestModel(pattern="latency", drive="local-ssd")
+    dumped = model.model_dump()
+
+    # model_dump() returns enum objects by default
+    assert dumped["pattern"] == AccessPattern.latency
+    assert dumped["drive"] == DriveType.local_ssd
+
+    # But they should compare equal to strings because they ARE strings
+    assert dumped["pattern"] == "latency"
+    assert dumped["drive"] == "local-ssd"
+
+    # And isinstance should work
+    assert isinstance(dumped["pattern"], str)
+    assert isinstance(dumped["drive"], str)

service_capacity_modeling-0.3.88/tests/test_enum_utils.py

@@ -1,140 +0,0 @@
-import pytest
-from pydantic import BaseModel
-
-from service_capacity_modeling.interface import AccessConsistency
-from service_capacity_modeling.interface import AccessPattern
-from service_capacity_modeling.interface import BufferComponent
-from service_capacity_modeling.interface import BufferIntent
-from service_capacity_modeling.interface import DriveType
-from service_capacity_modeling.interface import IntervalModel
-from service_capacity_modeling.interface import Platform
-
-# List of all enums that should have per-member docstrings
-DOCUMENTED_ENUMS = [
-    IntervalModel,
-    DriveType,
-    Platform,
-    AccessPattern,
-    AccessConsistency,
-    BufferComponent,
-    BufferIntent,
-]
-
-
-@pytest.mark.parametrize("enum_class", DOCUMENTED_ENUMS)
-def test_enums_have_docstrings(enum_class):
-    """Test that all interface.py enums have comprehensive per-member
-    docstrings
-
-    This test ensures that all enum members have their own
-    runtime-accessible docstrings, which makes them discoverable via
-    help(), IDE tooltips, and the __doc__ attribute.
-
-    The enums use the @enum_docstrings decorator which parses source code
-    to attach docstrings that appear below each member (following PEP 257
-    attribute docstring conventions).
-
-    See: https://stackoverflow.com/questions/19330460/
-    how-do-i-put-docstrings-on-enums
-    """
-    enum_name = enum_class.__name__
-
-    # Check class has a docstring
-    assert enum_class.__doc__ is not None, f"{enum_name} must have a class docstring"
-    assert len(enum_class.__doc__) > 0, f"{enum_name} class docstring must not be empty"
-
-    # Check each member has its own unique docstring
-    for member in enum_class:
-        assert member.__doc__ is not None, (
-            f"{enum_name}.{member.name} must have a docstring. "
-            f"Add a docstring after the member definition:\n"
-            f'    {member.name} = "{member.value}"\n'
-            f'    """Your documentation here"""'
-        )
-        assert len(member.__doc__.strip()) > 0, (
-            f"{enum_name}.{member.name} docstring must not be empty"
-        )
-        # Verify it's not just the class docstring (should be member-specific)
-        assert member.__doc__ != enum_class.__doc__, (
-            f"{enum_name}.{member.name} should have its own docstring, "
-            f"not inherit the class docstring. "
-            f"Did the @enum_docstrings decorator work?"
-        )
-
-    # Verify different members have different docstrings
-    members = list(enum_class)
-    if len(members) >= 2:
-        assert members[0].__doc__ != members[1].__doc__, (
-            f"{enum_name}: Different enum members should have different docstrings"
-        )
-
-
-@pytest.mark.parametrize("enum_class", DOCUMENTED_ENUMS)
-def test_enums_json_schema_includes_member_docstrings(enum_class):
-    """Test that enum member docstrings appear in Pydantic JSON schemas
-
-    The @enum_docstrings decorator adds __get_pydantic_json_schema__ to
-    generate oneOf schemas with per-member descriptions. This ensures
-    enum documentation is available in API schemas, OpenAPI specs, etc.
-    """
-    enum_name = enum_class.__name__
-
-    # Create a test model using this enum
-    TestModel = type(
-        "TestModel",
-        (BaseModel,),
-        {"__annotations__": {"field": enum_class}},
-    )
-
-    # Get the JSON schema
-    schema = TestModel.model_json_schema()
-
-    # Check the enum definition exists in $defs
-    assert "$defs" in schema, f"{enum_name}: Schema missing $defs"
-    assert enum_name in schema["$defs"], f"{enum_name}: Enum not in $defs"
-
-    enum_schema = schema["$defs"][enum_name]
-
-    # Check oneOf exists with member descriptions
-    assert "oneOf" in enum_schema, (
-        f"{enum_name}: JSON schema missing oneOf for member descriptions"
-    )
-
-    one_of = enum_schema["oneOf"]
-    assert len(one_of) == len(enum_class), (
-        f"{enum_name}: oneOf should have {len(enum_class)} entries"
-    )
-
-    # Verify each member has proper schema entry
-    for member in enum_class:
-        matching_entries = [
-            entry for entry in one_of if entry.get("const") == member.value
-        ]
-
-        assert len(matching_entries) == 1, (
-            f"{enum_name}.{member.name}: Should have exactly one oneOf entry"
-        )
-
-        entry = matching_entries[0]
-
-        # Check required fields
-        assert "const" in entry, f"{enum_name}.{member.name}: Missing 'const'"
-        assert "title" in entry, f"{enum_name}.{member.name}: Missing 'title'"
-        assert "description" in entry, (
-            f"{enum_name}.{member.name}: Missing 'description'"
-        )
-
-        # Verify description matches member docstring
-        assert entry["description"] == member.__doc__, (
-            f"{enum_name}.{member.name}: Schema description doesn't match "
-            f"member.__doc__"
-        )
-
-        # Verify description is not empty and not the class docstring
-        assert len(entry["description"].strip()) > 0, (
-            f"{enum_name}.{member.name}: Description should not be empty"
-        )
-        assert entry["description"] != enum_class.__doc__, (
-            f"{enum_name}.{member.name}: Description should be member-specific, "
-            f"not the class docstring"
-        )