Gergő Jedlicska
106 lines
3.1 KiB
Python
106 lines
3.1 KiB
Python
from dataclasses import dataclass
|
|
from typing import Dict, List, Optional
|
|
from unittest import TestCase
|
|
|
|
from specklepy.objects.base import Base
|
|
from specklepy.objects.graph_traversal.traversal import GraphTraversal, TraversalRule
|
|
|
|
|
|
@dataclass()
|
|
class TraversalMock(Base):
|
|
child: Optional[Base]
|
|
list_children: List[Base]
|
|
dict_children: Dict[str, Base]
|
|
|
|
|
|
class GraphTraversalTests(TestCase):
|
|
def test_traverse_list_members(self):
|
|
traverse_lists_rule = TraversalRule(
|
|
[lambda _: True],
|
|
lambda x: [
|
|
item
|
|
for item in x.get_member_names()
|
|
if isinstance(getattr(x, item, None), list)
|
|
],
|
|
)
|
|
|
|
expected_traverse = Base()
|
|
expected_traverse.id = "List Member"
|
|
|
|
expected_ignore = Base()
|
|
expected_ignore.id = "Not List Member"
|
|
|
|
test_case = TraversalMock(
|
|
list_children=[expected_traverse],
|
|
dict_children={"myprop": expected_ignore},
|
|
child=expected_ignore,
|
|
)
|
|
|
|
ret = [
|
|
context.current
|
|
for context in GraphTraversal([traverse_lists_rule]).traverse(test_case)
|
|
]
|
|
|
|
self.assertCountEqual(ret, [test_case, expected_traverse])
|
|
self.assertNotIn(expected_ignore, ret)
|
|
self.assertEqual(len(ret), 2)
|
|
|
|
def test_traverse_dict_members(self):
|
|
traverse_lists_rule = TraversalRule(
|
|
[lambda _: True],
|
|
lambda x: [
|
|
item
|
|
for item in x.get_member_names()
|
|
if isinstance(getattr(x, item, None), dict)
|
|
],
|
|
)
|
|
|
|
expected_traverse = Base()
|
|
expected_traverse.id = "Dict Member"
|
|
|
|
expected_ignore = Base()
|
|
expected_ignore.id = "Not Dict Member"
|
|
|
|
test_case = TraversalMock(
|
|
list_children=[expected_ignore],
|
|
dict_children={"myprop": expected_traverse},
|
|
child=expected_ignore,
|
|
)
|
|
|
|
ret = [
|
|
context.current
|
|
for context in GraphTraversal([traverse_lists_rule]).traverse(test_case)
|
|
]
|
|
|
|
self.assertCountEqual(ret, [test_case, expected_traverse])
|
|
self.assertNotIn(expected_ignore, ret)
|
|
self.assertEqual(len(ret), 2)
|
|
|
|
def test_traverse_dynamic(self):
|
|
traverse_lists_rule = TraversalRule(
|
|
[lambda _: True], lambda x: x.get_dynamic_member_names()
|
|
)
|
|
|
|
expected_traverse = Base()
|
|
expected_traverse.id = "List Member"
|
|
|
|
expected_ignore = Base()
|
|
expected_ignore.id = "Not List Member"
|
|
|
|
test_case = TraversalMock(
|
|
child=expected_ignore,
|
|
list_children=[expected_ignore],
|
|
dict_children={"myprop": expected_ignore},
|
|
)
|
|
test_case["dynamicChild"] = expected_traverse
|
|
test_case["dynamicListChild"] = [expected_traverse]
|
|
|
|
ret = [
|
|
context.current
|
|
for context in GraphTraversal([traverse_lists_rule]).traverse(test_case)
|
|
]
|
|
|
|
self.assertCountEqual(ret, [test_case, expected_traverse, expected_traverse])
|
|
self.assertNotIn(expected_ignore, ret)
|
|
self.assertEqual(len(ret), 3)
|