Add more assertions on dep tree

Author: donatasm

test/unit/test_import.py

@@ -10,8 +10,11 @@ def test_import_layer_does_not_import_heavy_libs():
     _, raw_output = get_import_time(import_cmd="layer", module_only=True)
     all_imports = prune_import_depth(parse_import_time(raw_output), depth=4)
 
-    deps = import_tree_to_waterfall(all_imports, width=64)
+    assert len(all_imports) > 0
+    assert all_imports[0].name == "root"
+    assert "layer" in {nested.name for nested in all_imports[0].nested_imports}
 
+    deps = import_tree_to_waterfall(all_imports, width=64)
     assert "torch" not in deps
     assert "xgboost" not in deps
     assert "keras" not in deps