API output assert with types

Author: priyankabose

Diff for: slither/analyses/data_dependency/data_dependency.py

@@ -293,7 +293,7 @@ def get_all_dependencies_ssa(
     return context.context[KEY_SSA]
 
 
-def get_must_depends_on(variable: SUPPORTED_TYPES) -> SUPPORTED_TYPES | None:
+def get_must_depends_on(variable: SUPPORTED_TYPES) -> List:
     """
     Return must dependency of a variable if exist otherwise return None.
 
@@ -305,7 +305,8 @@ def get_must_depends_on(variable: SUPPORTED_TYPES) -> SUPPORTED_TYPES | None:
         return []
     return [list(must_dependencies)[0]]
 
-def compute_must_dependencies(v:SUPPORTED_TYPES) -> Set[Variable]:
+
+def compute_must_dependencies(v: SUPPORTED_TYPES) -> Set[Variable]:
     if isinstance(v, (SolidityVariableComposed, Constant)) or (
         v.function.visibility in ["public", "external"] and v in v.function.parameters
     ):
@@ -359,6 +360,7 @@ def compute_must_dependencies(v:SUPPORTED_TYPES) -> Set[Variable]:
             must_dependencies = must_dependencies.union(result)
     return must_dependencies
 
+
 # endregion
 ###################################################################################
 ###################################################################################

Diff for: tests/unit/core/test_must_depend_on.py

@@ -1,14 +1,23 @@
 from pathlib import Path
 from slither import Slither
-from slither.analyses.data_dependency.data_dependency import (
-    get_must_depends_on
+from slither.analyses.data_dependency.data_dependency import get_must_depends_on
+from slither.core.variables.variable import Variable
+from slither.core.declarations import (
+    SolidityVariable,
+)
+from typing import Union
+from slither.slithir.variables import (
+    Constant,
 )
 
 TEST_DATA_DIR = Path(__file__).resolve().parent / "test_data"
+SUPPORTED_TYPES = Union[Variable, SolidityVariable, Constant]
+
 
 def test_must_depend_on_returns(solc_binary_path):
     solc_path = solc_binary_path("0.8.19")
     file = Path(TEST_DATA_DIR, "must_depend_on.sol").as_posix()
     slither_obj = Slither(file, solc=solc_path)
     result = get_must_depends_on(slither_obj.contracts[1].functions[2].parameters[0])
-    assert isinstance(result, list) and len(result) <= 1
+    assert isinstance(result, list)
+    assert len(result) == 0 or (len(result) == 1 and result[0] in SUPPORTED_TYPES)