runs lint for test_dataset_builder

Author: patrickmcarlos

src/sagemaker/feature_store/dataset_builder.py

@@ -299,16 +299,16 @@ def with_feature_group(
                 will be used as a join key (default: None).
             included_feature_names (List[str]): A list of strings representing features to be
                 included in the output (default: None).
-        feature_name_in_target (str): A string representing the feature in the target feature group
-            that will be compared to the target feature in the base feature group. If None is
-            provided, the record identifier will be used in the join statement. (default: None).
-        join_comparator (JoinComparatorEnum): A JoinComparatorEnum representing the comparator used
-            when joining the target feature in the base feature group and the feature in the target
-            feature group (default: None).
-        join_type (JoinTypeEnum): A JoinTypeEnum representing the type of join between the base and
-            target feature groups. (default: None).
-        Returns:
-            This DatasetBuilder object.
+            feature_name_in_target (str): A string representing the feature in the target feature group
+                that will be compared to the target feature in the base feature group. If None is
+                provided, the record identifier will be used in the join statement. (default: None).
+            join_comparator (JoinComparatorEnum): A JoinComparatorEnum representing the comparator used
+                when joining the target feature in the base feature group and the feature in the target
+                feature group (default: None).
+            join_type (JoinTypeEnum): A JoinTypeEnum representing the type of join between the base and
+                target feature groups. (default: None).
+            Returns:
+                This DatasetBuilder object.
         """
         self._feature_groups_to_be_merged.append(
             construct_feature_group_to_be_merged(

tests/unit/sagemaker/feature_store/test_dataset_builder.py

@@ -163,9 +163,15 @@ def test_with_feature_group_with_additional_params(sagemaker_session_mock):
             {"FeatureName": "feature-2", "FeatureType": "String"},
         ],
     }
-    dataset_builder.with_feature_group(feature_group, "target-feature", [
-                                       "feature-1", "feature-2"], join_comparator=JoinComparatorEnum.LESS_THAN, join_type=JoinTypeEnum.LEFT_JOIN, feature_name_in_target='feature-2')
-                                       
+    dataset_builder.with_feature_group(
+        feature_group,
+        "target-feature",
+        ["feature-1", "feature-2"],
+        join_comparator=JoinComparatorEnum.LESS_THAN,
+        join_type=JoinTypeEnum.LEFT_JOIN,
+        feature_name_in_target="feature-2",
+    )
+
     assert len(dataset_builder._feature_groups_to_be_merged) == 1
     assert dataset_builder._feature_groups_to_be_merged[0].features == [
         "feature-1",
@@ -189,10 +195,7 @@ def test_with_feature_group_with_additional_params(sagemaker_session_mock):
         dataset_builder._feature_groups_to_be_merged[0].event_time_identifier_feature.feature_type
         == FeatureTypeEnum.STRING
     )
-    assert (
-        dataset_builder._feature_groups_to_be_merged[0].join_type
-        == JoinTypeEnum.LEFT_JOIN
-    )
+    assert dataset_builder._feature_groups_to_be_merged[0].join_type == JoinTypeEnum.LEFT_JOIN
     assert (
         dataset_builder._feature_groups_to_be_merged[0].join_comparator
         == JoinComparatorEnum.LESS_THAN