diff --git a/metadata-ingestion/src/datahub/utilities/memory_footprint.py b/metadata-ingestion/src/datahub/utilities/memory_footprint.py
index 9fc1900441..fa9e64cbf4 100644
--- a/metadata-ingestion/src/datahub/utilities/memory_footprint.py
+++ b/metadata-ingestion/src/datahub/utilities/memory_footprint.py
@@ -1,7 +1,7 @@
 from collections import deque
 from itertools import chain
 from sys import getsizeof
-from typing import Any, Dict
+from typing import Any, Callable
 
 
 def total_size(o: Any, handlers: Any = {}) -> int:
@@ -15,8 +15,7 @@ def total_size(o: Any, handlers: Any = {}) -> int:
     Based on https://github.com/ActiveState/recipe-577504-compute-mem-footprint/blob/master/recipe.py
     """
 
-    def dict_handler(d: Dict) -> chain[Any]:
-        return chain.from_iterable(d.items())
+    dict_handler: Callable[[Any], chain[Any]] = lambda d: chain.from_iterable(d.items())
 
     all_handlers = {
         tuple: iter,
diff --git a/metadata-ingestion/tests/unit/utilities/test_memory_footprint.py b/metadata-ingestion/tests/unit/utilities/test_memory_footprint.py
new file mode 100644
index 0000000000..703cb55f55
--- /dev/null
+++ b/metadata-ingestion/tests/unit/utilities/test_memory_footprint.py
@@ -0,0 +1,27 @@
+from collections import defaultdict
+
+from datahub.utilities import memory_footprint
+
+
+def test_total_size_with_empty_dict():
+    size = memory_footprint.total_size({})
+    # Only asserting if it is bigger than 0 because the actual sizes differs per python version
+    assert size > 0
+
+
+def test_total_size_with_list():
+    size = memory_footprint.total_size({"1": [1, 2, 3, 4]})
+    # Only asserting if it is bigger than 0 because the actual sizes differs per python version
+    assert size > 0
+
+
+def test_total_size_with_none():
+    size = memory_footprint.total_size(None)
+    # Only asserting if it is bigger than 0 because the actual sizes differs per python version
+    assert size > 0
+
+
+def test_total_size_with_defaultdict():
+    size = memory_footprint.total_size(defaultdict)
+    # Only asserting if it is bigger than 0 because the actual sizes differs per python version
+    assert size > 0