From c61a7098fd3b21b76f60edac1690da9648d41b97 Mon Sep 17 00:00:00 2001
From: Varun Mittal <varunmittal91@gmail.com>
Date: Mon, 18 Dec 2023 08:07:06 -0500
Subject: [PATCH] Fixed test cases for check_type_name.

Signed-off-by: Varun Mittal <varunmittal91@gmail.com>
---
 focus_validator/config_objects/rule.py        | 42 ++++++++-----------
 .../test_check_type_friendly_name.py          |  4 ++
 2 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/focus_validator/config_objects/rule.py b/focus_validator/config_objects/rule.py
index 47a2a61..abac07a 100644
--- 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
index a5fb610..1395eec 100644
--- 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):