test: different call method of viz_folium

src/yuheng/plugin/folium/__main__.py → src/yuheng/plugin/viz_folium/__main__.py

@@ -11,7 +11,7 @@
 from yuheng.type import Node, Relation, Way
 
 
-class Map:
+class VizFolium:
     def __init__(self):
         self.element_list = []
 
@@ -22,9 +22,11 @@ def add(self, element: Union[Waifu, Way, Node, Relation]) -> None:
         self.element_list.append(element)
         return None
 
-    def display(self) -> None:
+    def display(self, **kwargs) -> None:
         """
         This func display some element to html
+
+        If **kwargs is not blank, then we should call self.add() for each one, then display as expected.
         """
 
         m = folium.Map(location=[11.4, 51.4], zoom_start=0)

tests/cases/plugin_folium.py → tests/cases/plugin_viz_folium.py

@@ -8,13 +8,11 @@
 
 from yuheng import Waifu
 from yuheng.type import Member, Node, Relation, Way
-from yuheng.plugin.folium.__main__ import display
+from yuheng.plugin.folium.__main__ import VizFolium, display
 
 
 class TestPluginVisualizationFolium(unittest.TestCase):
     def setUp(self) -> None:
-        from yuheng.plugin.folium.__main__ import display
-
         self.test_node_1 = Node(
             {"id": "1"}, {"name": "folium cafe", "amenity": "cafe"}
         )
@@ -34,9 +32,34 @@ def setUp(self) -> None:
             ],
         )
         self.test_map = Waifu()
+        # 需要插入到Waifu对象里面。
+        # self.test_map.insert(
+        #     self.test_node_1,
+        #     self.test_node_2,
+        #     self.test_way,
+        #     self.test_relation,
+        # )
+
+    def test_plugin_viz_folium_display_kwargs(self):
+        carto_viz = VizFolium()
+        carto_viz.display(
+            self.test_node_1,
+            self.test_node_2,
+            self.test_way,
+            self.test_relation,
+        )
+
+    def test_plugin_viz_folium_display_added(self):
+        carto_viz = VizFolium()
+        carto_viz.add(self.test_node_1)
+        carto_viz.add(self.test_node_2)
+        carto_viz.add(self.test_way)
+        carto_viz.add(self.test_relation)
+        carto_viz.display()
 
-    def test_plugin_driver_poly_import(self):
-        display(self.test_relation)
+    def test_plugin_viz_folium_display_object(self):
+        carto_viz = VizFolium()
+        carto_viz.display(self.test_map)
 
 
 if __name__ == "__main__":