Fixed test cases for check_type_name.

Signed-off-by: Varun Mittal <[email protected]>
finopsfoundation · Dec 18, 2023 · c61a709 · c61a709
1 parent 6d0b595
commit c61a709
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 24 deletions.
diff --git a/focus_validator/config_objects/rule.py b/focus_validator/config_objects/rule.py
@@ -37,36 +37,15 @@ class Rule(BaseModel):
     check_friendly_name: Annotated[
         Optional[str], Field(validate_default=True)
     ] = None  # auto generated or else can be overwritten
-    check_type_friendly_name: Optional[str] = None
+    check_type_friendly_name: Annotated[
+        Optional[str], Field(validate_default=True)
+    ] = None
 
     model_config = ConfigDict(
         extra="forbid",  # prevents config from containing any undesirable keys
         frozen=True,  # prevents any modification to any attribute onces loaded from config
     )
 
-    @model_validator(mode="before")
-    @classmethod
-    def root_val(cls, values):
-        """
-        Root validator that checks for all options passed in the config and generate missing options.
-        """
-        if values is None:
-            values = {}
-
-        check = values.get("check")
-        check_friendly_name = values.get("check_friendly_name")
-        column_id = values.get("column_id")
-        if check is not None:
-            if isinstance(check, str):
-                check_type_friendly_name = "".join(
-                    [word.title() for word in check.split("_")]
-                )
-            else:
-                check_type_friendly_name = check.__class__.__name__
-            values["check_type_friendly_name"] = check_type_friendly_name
-
-        return values
-
     @field_validator("check_friendly_name")
     def validate_or_generate_check_friendly_name(
         cls, check_friendly_name, validation_info: ValidationInfo
@@ -82,6 +61,21 @@ def validate_or_generate_check_friendly_name(
             )
         return check_friendly_name
 
+    @field_validator("check_type_friendly_name")
+    def validate_or_generate_check_type_friendly_name(
+        cls, check_type_friendly_name, validation_info: ValidationInfo
+    ):
+        values = validation_info.data
+        if values.get("check") is not None and values.get("column_id") is not None:
+            check = values.get("check")
+            if isinstance(check, str):
+                check_type_friendly_name = "".join(
+                    [word.title() for word in check.split("_")]
+                )
+            else:
+                check_type_friendly_name = check.__class__.__name__
+        return check_type_friendly_name
+
     @staticmethod
     def load_yaml(
         rule_path, column_namespace: Optional[str] = None

diff --git a/tests/config_objects/test_check_type_friendly_name.py b/tests/config_objects/test_check_type_friendly_name.py
@@ -49,11 +49,15 @@ def test_random_value_is_ignored(self):
         self.assertEqual(sample.check_type_friendly_name, "CheckUnique")
 
     def test_data_type_config(self):
+        # Ensures that the check_type_friendly_name is generated correctly for DataTypeCheck
+
         model_factory = ModelFactory.create_factory(model=Rule)
 
+        # generate random rule object
         sample_data_type = model_factory.build(
             **{"check": DataTypeCheck(data_type=DataTypes.STRING)}
         )
+
         self.assertEqual(sample_data_type.check_type_friendly_name, "DataTypeCheck")
 
     def test_check_type_config_deny_update(self):