From 68160cb53d40b93ba6a4afa77f18f62a3cfc5bc2 Mon Sep 17 00:00:00 2001
From: Zijian Zhang <doomspec@outlook.com>
Date: Mon, 9 Oct 2023 03:36:21 -0400
Subject: [PATCH] test: add test case for .tree.yml

---
 doc_in_py/core.py                   |  3 ++-
 doc_in_py/test/test_module_multi.py | 25 +++++++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)
 create mode 100644 doc_in_py/test/test_module_multi.py

diff --git a/doc_in_py/core.py b/doc_in_py/core.py
index c0eaf77..f1355a6 100644
--- a/doc_in_py/core.py
+++ b/doc_in_py/core.py
@@ -49,6 +49,7 @@ def extract_module_tree_without_comment(module, root_path):
     # Get sub_module tree from .tree.yml if it exists
     module_dir = os.path.dirname(inspect.getfile(module))
     tree_yml_path = os.path.join(module_dir, ".tree.yml")
+    tree_config_dict = {}
     if os.path.exists(tree_yml_path):
         try:
             with open(tree_yml_path, "r") as f:
@@ -56,7 +57,7 @@ def extract_module_tree_without_comment(module, root_path):
         except:
             pass
 
-    if module.__file__.endswith("__init__.py"):
+    if module.__file__.endswith("__init__.py") and len(tree_config_dict) > 0:
         sections_dict = build_module_section_dict(tree_config_dict, sub_modules)
         if sections_dict is not None:
             sub_modules = sections_dict
diff --git a/doc_in_py/test/test_module_multi.py b/doc_in_py/test/test_module_multi.py
new file mode 100644
index 0000000..0d1c02e
--- /dev/null
+++ b/doc_in_py/test/test_module_multi.py
@@ -0,0 +1,25 @@
+from doc_in_py.core import get_module_members
+from utils import to_dict
+
+
+def test_section():
+    import testing_module_multi
+    struct = get_module_members(testing_module_multi)
+    actual = to_dict(struct)
+    expected = {'struct_type': 'module', 'name': 'testing_module_multi', 'children': [
+        {'struct_type': 'section', 'name': 'first section', 'children': [
+            {'struct_type': 'comment',
+             'obj': 'This is some annotation to the first section'},
+            {'struct_type': 'module', 'name': 'testing_module_multi.a',
+             'children': [{'struct_type': 'class', 'name': 'A'}]},
+            {'struct_type': 'module', 'name': 'testing_module_multi.b'}]},
+        {'struct_type': 'section', 'name': 'second section',
+         'children': [{'struct_type': 'module', 'name': 'testing_module_multi.c'},
+                      {'struct_type': 'section', 'name': 'another section', 'children': [
+                          {'struct_type': 'comment', 'obj': 'This is another section'},
+                          {'struct_type': 'module', 'name': 'testing_module_multi.d'},
+                          {'struct_type': 'module', 'name': 'testing_module_multi.e'}]}]},
+        {'struct_type': 'section', 'name': 'other section', 'children': [
+            {'struct_type': 'module', 'name': 'testing_module_multi.another'}]}]}
+
+    assert actual == expected