moved fun from task_utils to utils

tests/integration/tasks/test_task_utils.py

@@ -3,7 +3,7 @@
 from prefect.engine.state import Failed, Success
 from prefect.tasks.secrets import PrefectSecret
 
-from viadot.task_utils import custom_mail_state_handler, set_new_kv, check_value
+from viadot.task_utils import custom_mail_state_handler, set_new_kv
 
 
 def test_custom_state_handler():
@@ -27,53 +27,4 @@ def test_set_new_kv():
     set_new_kv.run(kv_name="test_for_setting_kv", df=df, filter_column="col1")
     result = get_key_value("test_for_setting_kv")
     assert result == "72"
-    set_key_value(key="test_for_setting_kv", value=None)
-
-
-# Sample test checking the correctness of the function when the key is found
-def test_check_value_found():
-    json_data = {
-        "first_known_lvl": {
-            "second_known_lvl": {
-                "third_known_lvl": {
-                    "searched_phrase": "phrase"
-                }
-            }
-        }
-    }
-    result = check_value(json_data["first_known_lvl"]["second_known_lvl"]["third_known_lvl"], ["searched_phrase"])
-    assert result == "phrase"
-
-# Sample test checking the correctness of the function when the key is not found
-def test_check_value_not_found():
-    json_data = {
-        "first_known_lvl": {
-            "second_known_lvl": {
-                "third_known_lvl": {
-                    "other_phrase": "This won't be found"
-                }
-            }
-        }
-    }
-    result = check_value(json_data["first_known_lvl"]["second_known_lvl"]["third_known_lvl"], ["searched_phrase"])
-    assert result is None
-
-# Sample test checking the correctness of the function with an empty dictionary
-def test_check_value_empty_dict():
-    json_data = {}
-    result = check_value(json_data, ["searched_phrase"])
-    assert result is None
-
-# Sample test checking the correctness of the function with a nonexistent key
-def test_check_value_nonexistent_key():
-    json_data = {
-        "first_known_lvl": {
-            "second_known_lvl": {
-                "third_known_lvl": {
-                    "searched_phrase": "phrase"
-                }
-            }
-        }
-    }
-    result = check_value(json_data, ["nonexistent_key"])
-    assert result is None
+    set_key_value(key="test_for_setting_kv", value=None)

tests/unit/test_utils.py

@@ -9,6 +9,7 @@
     add_viadot_metadata_columns,
     check_if_empty_file,
     gen_bulk_insert_query_from_df,
+    check_value,
 )
 
 EMPTY_CSV_PATH = "empty.csv"
@@ -153,3 +154,51 @@ def test_add_viadot_metadata_columns_with_parameter():
     assert df_base.columns.to_list() == ["a", "b"]
     assert df_decorated.columns.to_list() == ["a", "b", "_viadot_source"]
     assert df_decorated["_viadot_source"][0] == "Source_name"
+
+# Sample test checking the correctness of the function when the key is found
+def test_check_value_found():
+    json_data = {
+        "first_known_lvl": {
+            "second_known_lvl": {
+                "third_known_lvl": {
+                    "searched_phrase": "phrase"
+                }
+            }
+        }
+    }
+    result = check_value(json_data["first_known_lvl"]["second_known_lvl"]["third_known_lvl"], ["searched_phrase"])
+    assert result == "phrase"
+
+# Sample test checking the correctness of the function when the key is not found
+def test_check_value_not_found():
+    json_data = {
+        "first_known_lvl": {
+            "second_known_lvl": {
+                "third_known_lvl": {
+                    "other_phrase": "This won't be found"
+                }
+            }
+        }
+    }
+    result = check_value(json_data["first_known_lvl"]["second_known_lvl"]["third_known_lvl"], ["searched_phrase"])
+    assert result is None
+
+# Sample test checking the correctness of the function with an empty dictionary
+def test_check_value_empty_dict():
+    json_data = {}
+    result = check_value(json_data, ["searched_phrase"])
+    assert result is None
+
+# Sample test checking the correctness of the function with a nonexistent key
+def test_check_value_nonexistent_key():
+    json_data = {
+        "first_known_lvl": {
+            "second_known_lvl": {
+                "third_known_lvl": {
+                    "searched_phrase": "phrase"
+                }
+            }
+        }
+    }
+    result = check_value(json_data, ["nonexistent_key"])
+    assert result is None

viadot/task_utils.py

@@ -791,29 +791,4 @@ def validate_df(df: pd.DataFrame, tests: dict = None) -> None:
         failed_tests_msg = ", ".join(failed_tests_list)
         raise ValidationError(
             f"Validation failed for {failed_tests} test/tests: {failed_tests_msg}"
-        )
-
-
-def check_value(base, lvls: List):
-    """
-            Task to extract data from nested json file if there is any under passed parameters.
-            Otherwise return None.
-
-    Args:
-        base: variable with base lvl of the json, fo example:
-                json_file["first_known_lvl"]["second_known_lvl"]["third_known_lvl"]
-        lvls (List): List of potential lower levels of nested json for data retrieval. For example:
-                ["first_lvl_below_base", "second_lvl_below_base", "searched_phrase"]
-
-    Return:
-        Searched value for the lowest level, in example data under "searched_phrase" key.
-    """
-
-    for lvl in lvls:
-        if isinstance(base, dict):
-            base = base.get(lvl)
-            if base is None:
-                return None
-        else:
-            return base
-    return base
+        )

viadot/tasks/genesys.py

@@ -13,6 +13,7 @@
 
 from viadot.exceptions import APIError
 from viadot.sources import Genesys
+from viadot.utils import check_value
 from viadot.task_utils import *
 
 logger = logging.get_logger()

viadot/utils.py

@@ -2,7 +2,7 @@
 import os
 import re
 from itertools import chain
-from typing import Any, Callable, Dict, List, Literal
+from typing import Union, Any, Callable, Dict, List, Literal
 
 import pandas as pd
 import prefect
@@ -460,3 +460,28 @@ def get_nested_dict(d):
                 return d
     else:
         return None
+
+
+def check_value(base: Union[Dict, Any], levels: List) -> Union[None, Any]:
+    """        
+    Task to extract data from nested json file if there is any under passed parameters.
+    Otherwise return None.
+
+    Args:
+        base (Dict, Any): variable with base lvl of the json, for example:
+                          json_file["first_known_lvl"]["second_known_lvl"]["third_known_lvl"]
+        levels (List): List of potential lower levels of nested json for data retrieval. For example:
+                       ["first_lvl_below_base", "second_lvl_below_base", "searched_phrase"]
+
+    Returns:
+        Union[None, Any]: Searched value for the lowest level, in example data under "searched_phrase" key.
+    """
+
+    for lvl in levels:
+        if isinstance(base, dict):
+            base = base.get(lvl)
+            if base is None:
+                return None
+        else:
+            return base
+    return base