fix(feature-store): Fix FeatureGroupManager code issues and improve test coverage

- Use isinstance() for Unassigned checks instead of == Unassigned()
- Add class-level type annotation for _lf_client_cache
- Replace fragile docstring inheritance with proper docstring
- Fix create() to return FeatureGroupManager instead of FeatureGroup
  by calling cls.get() after super().create()
- Update create() return type annotation to Optional[FeatureGroupManager]
- Add feature_group_arn validation before S3 policy generation
- Fix integ test logger name (feature_group -> feature_group_manager)
- Rename test_lakeformation.py to test_feature_group_manager.py
- Add unit tests for: return type verification, Iceberg table format
  S3 path handling, missing ARN validation, happy-path return values,
  session/region pass-through, and region inference from session
---
X-AI-Prompt: Review FeatureGroupManager class, fix identified issues, increase test coverage
X-AI-Tool: kiro-cli

Author: BassemHalim

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_group_manager.py

@@ -56,10 +56,9 @@ class LakeFormationConfig(Base):
 
 
 class FeatureGroupManager(FeatureGroup):
+    """FeatureGroup with extended management capabilities."""
 
-    # Inherit parent docstring and append our additions
-    if FeatureGroup.__doc__ and __doc__:
-        __doc__ = FeatureGroup.__doc__
+    _lf_client_cache: dict
 
     @staticmethod
     def _s3_uri_to_arn(s3_uri: str, region: Optional[str] = None) -> str:
@@ -221,8 +220,8 @@ def _get_lake_formation_client(
             A boto3 Lake Formation client.
         """
         cache_key = (id(session), region)
-        if not hasattr(self, "_lf_client_cache"):
-            self._lf_client_cache: dict = {}
+        if not hasattr(self, "_lf_client_cache") or self._lf_client_cache is None:
+            self._lf_client_cache = {}
 
         if cache_key not in self._lf_client_cache:
             boto_session = session or Session()
@@ -470,14 +469,14 @@ def enable_lake_formation(
             )
 
         # Validate offline store exists
-        if self.offline_store_config is None or self.offline_store_config == Unassigned():
+        if self.offline_store_config is None or isinstance(self.offline_store_config, Unassigned):
             raise ValueError(
                 f"Feature Group '{self.feature_group_name}' does not have an offline store configured. "
                 "Lake Formation can only be enabled for Feature Groups with offline stores."
             )
 
         # Get role ARN from Feature Group config
-        if self.role_arn is None or self.role_arn == Unassigned():
+        if self.role_arn is None or isinstance(self.role_arn, Unassigned):
             raise ValueError(
                 f"Feature Group '{self.feature_group_name}' does not have a role_arn configured. "
                 "Lake Formation requires a role ARN to grant permissions."
@@ -493,7 +492,7 @@ def enable_lake_formation(
             raise ValueError("Offline store S3 configuration is missing")
 
         resolved_s3_uri = s3_config.resolved_output_s3_uri
-        if resolved_s3_uri is None or resolved_s3_uri == Unassigned():
+        if resolved_s3_uri is None or isinstance(resolved_s3_uri, Unassigned):
             raise ValueError(
                 "Resolved S3 URI not available. Ensure the Feature Group is in 'Created' status."
             )
@@ -599,6 +598,13 @@ def enable_lake_formation(
 
         # Generate and optionally print S3 deny policy
         if show_s3_policy:
+            # Validate feature_group_arn is available for account ID extraction
+            if self.feature_group_arn is None or isinstance(self.feature_group_arn, Unassigned):
+                raise ValueError(
+                    "Feature Group ARN is required to generate S3 policy. "
+                    "Ensure the Feature Group is in 'Created' status."
+                )
+
             # Extract bucket name and prefix from resolved S3 URI using core utility
             bucket_name, s3_prefix = parse_s3_url(resolved_s3_uri_str)
 
@@ -659,9 +665,9 @@ def create(
         lake_formation_config: Optional[LakeFormationConfig] = None,
         session: Optional[Session] = None,
         region: Optional[StrPipeVar] = None,
-    ) -> Optional["FeatureGroup"]:
+    ) -> Optional["FeatureGroupManager"]:
         """
-        Create a FeatureGroup resource with optional Lake Formation governance.
+        Create a FeatureGroupManager resource with optional Lake Formation governance.
 
         Parameters:
             feature_group_name: The name of the FeatureGroup.
@@ -682,7 +688,7 @@ def create(
             region: Region name.
 
         Returns:
-            The FeatureGroup resource.
+            The FeatureGroupManager resource.
 
         Raises:
             botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
@@ -744,7 +750,10 @@ def create(
         if use_pre_prod_offline_store_replicator_lambda is not None:
             create_kwargs["use_pre_prod_offline_store_replicator_lambda"] = use_pre_prod_offline_store_replicator_lambda
 
-        feature_group = super().create(**create_kwargs)
+        super().create(**create_kwargs)
+
+        # Get as FeatureGroupManager instance (super().create() returns FeatureGroup)
+        feature_group = cls.get(feature_group_name=feature_group_name, session=session, region=region)
 
         # Enable Lake Formation if requested
         if lake_formation_config is not None and lake_formation_config.enabled:

sagemaker-mlops/tests/integ/test_featureStore_lakeformation.py

@@ -498,7 +498,7 @@ def test_enable_lake_formation_full_flow_with_policy_output(s3_uri, role, region
         assert fg.feature_group_status == "Created"
 
         # Enable Lake Formation governance with policy output
-        with caplog.at_level(logging.INFO, logger="sagemaker.mlops.feature_store.feature_group"):
+        with caplog.at_level(logging.INFO, logger="sagemaker.mlops.feature_store.feature_group_manager"):
             result = fg.enable_lake_formation(show_s3_policy=True)
 
         # Verify all phases completed successfully
@@ -573,7 +573,7 @@ def test_enable_lake_formation_no_policy_output_by_default(s3_uri, role, region,
         assert fg.feature_group_status == "Created"
 
         # Enable Lake Formation governance WITHOUT policy output (default)
-        with caplog.at_level(logging.INFO, logger="sagemaker.mlops.feature_store.feature_group"):
+        with caplog.at_level(logging.INFO, logger="sagemaker.mlops.feature_store.feature_group_manager"):
             result = fg.enable_lake_formation()
 
         # Verify all phases completed successfully
@@ -624,7 +624,7 @@ def test_enable_lake_formation_with_custom_role_policy_output(s3_uri, role, regi
 
         # Enable Lake Formation with custom registration role and policy output
         # Using the same role for both execution and registration for test simplicity
-        with caplog.at_level(logging.INFO, logger="sagemaker.mlops.feature_store.feature_group"):
+        with caplog.at_level(logging.INFO, logger="sagemaker.mlops.feature_store.feature_group_manager"):
             result = fg.enable_lake_formation(
                 use_service_linked_role=False,
                 registration_role_arn=role,