Skip to content

Commit 8fccc33

Browse files
jm-riverajm-rivera
committed
Add tests
1 parent 1cf83d4 commit 8fccc33

File tree

2 files changed

+445
-40
lines changed

2 files changed

+445
-40
lines changed

‎datacommons_client/tests/endpoints/test_node_endpoint.py

Lines changed: 181 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
from unittest.mock import MagicMock
2+
from unittest.mock import patch
23

34
from datacommons_client.endpoints.base import API
45
from datacommons_client.endpoints.node import NodeEndpoint
56
from datacommons_client.endpoints.response import NodeResponse
7+
from datacommons_client.models.graph import Parent
8+
from datacommons_client.models.node import Node
69

710

811
def test_node_endpoint_initialization():
@@ -30,13 +33,15 @@ def test_node_endpoint_fetch():
3033
endpoint = NodeEndpoint(api=api_mock)
3134
response = endpoint.fetch(node_dcids="test_node", expression="name")
3235

33-
api_mock.post.assert_called_once_with(payload={
34-
"nodes": ["test_node"],
35-
"property": "name"
36-
},
37-
endpoint="node",
38-
all_pages=True,
39-
next_token=None)
36+
api_mock.post.assert_called_once_with(
37+
payload={
38+
"nodes": ["test_node"],
39+
"property": "name"
40+
},
41+
endpoint="node",
42+
all_pages=True,
43+
next_token=None,
44+
)
4045
assert isinstance(response, NodeResponse)
4146
assert "test_node" in response.data
4247

@@ -80,13 +85,15 @@ def test_node_endpoint_fetch_property_values_out():
8085
out=True)
8186

8287
expected_expression = "->name{typeOf:City}"
83-
api_mock.post.assert_called_once_with(payload={
84-
"nodes": ["node1"],
85-
"property": expected_expression
86-
},
87-
endpoint="node",
88-
all_pages=True,
89-
next_token=None)
88+
api_mock.post.assert_called_once_with(
89+
payload={
90+
"nodes": ["node1"],
91+
"property": expected_expression
92+
},
93+
endpoint="node",
94+
all_pages=True,
95+
next_token=None,
96+
)
9097
assert isinstance(response, NodeResponse)
9198
assert "node1" in response.data
9299

@@ -112,13 +119,15 @@ def test_node_endpoint_fetch_property_values_in():
112119
out=False)
113120

114121
expected_expression = "<-name{typeOf:City}"
115-
api_mock.post.assert_called_once_with(payload={
116-
"nodes": ["node1"],
117-
"property": expected_expression
118-
},
119-
endpoint="node",
120-
all_pages=True,
121-
next_token=None)
122+
api_mock.post.assert_called_once_with(
123+
payload={
124+
"nodes": ["node1"],
125+
"property": expected_expression
126+
},
127+
endpoint="node",
128+
all_pages=True,
129+
next_token=None,
130+
)
122131
assert isinstance(response, NodeResponse)
123132
assert "node1" in response.data
124133

@@ -133,11 +142,13 @@ def test_node_endpoint_fetch_all_classes():
133142
}}))
134143

135144
response = endpoint.fetch_all_classes()
136-
endpoint.fetch_property_values.assert_called_once_with(node_dcids="Class",
137-
properties="typeOf",
138-
out=False,
139-
all_pages=True,
140-
next_token=None)
145+
endpoint.fetch_property_values.assert_called_once_with(
146+
node_dcids="Class",
147+
properties="typeOf",
148+
out=False,
149+
all_pages=True,
150+
next_token=None,
151+
)
141152
assert isinstance(response, NodeResponse)
142153
assert "Class" in response.data
143154

@@ -162,23 +173,153 @@ def test_node_endpoint_fetch_property_values_string_vs_list():
162173
properties="name",
163174
constraints=None,
164175
out=True)
165-
api_mock.post.assert_called_with(payload={
166-
"nodes": ["node1"],
167-
"property": "->name"
168-
},
169-
endpoint="node",
170-
all_pages=True,
171-
next_token=None)
176+
api_mock.post.assert_called_with(
177+
payload={
178+
"nodes": ["node1"],
179+
"property": "->name"
180+
},
181+
endpoint="node",
182+
all_pages=True,
183+
next_token=None,
184+
)
172185

173186
# List input
174187
response = endpoint.fetch_property_values(node_dcids="node1",
175188
properties=["name", "typeOf"],
176189
constraints=None,
177190
out=True)
178-
api_mock.post.assert_called_with(payload={
179-
"nodes": ["node1"],
180-
"property": "->[name, typeOf]"
181-
},
182-
endpoint="node",
183-
all_pages=True,
184-
next_token=None)
191+
api_mock.post.assert_called_with(
192+
payload={
193+
"nodes": ["node1"],
194+
"property": "->[name, typeOf]"
195+
},
196+
endpoint="node",
197+
all_pages=True,
198+
next_token=None,
199+
)
200+
201+
202+
@patch("datacommons_client.utils.graph.build_parents_dictionary")
203+
def test_fetch_entity_parents_as_dict(mock_build_parents_dict):
204+
"""Test fetch_entity_parents with dictionary output."""
205+
api_mock = MagicMock()
206+
endpoint = NodeEndpoint(api=api_mock)
207+
208+
api_mock.post.return_value = {
209+
"data": {
210+
"X": {
211+
"properties": {
212+
"containedInPlace": []
213+
}
214+
}
215+
}
216+
}
217+
endpoint.fetch_property_values = MagicMock()
218+
endpoint.fetch_property_values.return_value.get_properties.return_value = {
219+
"X": Node("X", "X name", types=["Country"])
220+
}
221+
222+
result = endpoint.fetch_entity_parents("X", as_dict=True)
223+
assert result == {"X": [{"dcid": "X", "name": "X name", "type": "Country"}]}
224+
225+
endpoint.fetch_property_values.assert_called_once_with(
226+
node_dcids="X", properties="containedInPlace")
227+
228+
229+
@patch("datacommons_client.utils.graph.build_parents_dictionary")
230+
def test_fetch_entity_parents_as_objects(mock_build_parents_dict):
231+
"""Test fetch_entity_parents with raw Parent object output."""
232+
api_mock = MagicMock()
233+
endpoint = NodeEndpoint(api=api_mock)
234+
235+
# Simulate what fetch_property_values().get_properties() would return
236+
endpoint.fetch_property_values = MagicMock()
237+
endpoint.fetch_property_values.return_value.get_properties.return_value = {
238+
"X": Node("X", "X name", types=["Country"])
239+
}
240+
241+
# Mock output of build_parents_dictionary
242+
parent_obj = Node("X", "X name", types=["Country"])
243+
mock_build_parents_dict.return_value = {"X": [parent_obj]}
244+
245+
result = endpoint.fetch_entity_parents("X", as_dict=False)
246+
247+
assert isinstance(result, dict)
248+
assert "X" in result
249+
assert isinstance(result["X"][0], Parent)
250+
251+
endpoint.fetch_property_values.assert_called_once_with(
252+
node_dcids="X", properties="containedInPlace")
253+
254+
255+
@patch("datacommons_client.endpoints.node.fetch_parents_lru")
256+
def test_fetch_parents_cached_delegates_to_lru(mock_fetch_lru):
257+
mock_fetch_lru.return_value = (Parent("B", "B name", "Region"),)
258+
endpoint = NodeEndpoint(api=MagicMock())
259+
260+
result = endpoint._fetch_parents_cached("X")
261+
262+
assert isinstance(result, tuple)
263+
assert result[0].dcid == "B"
264+
mock_fetch_lru.assert_called_once_with(endpoint, "X")
265+
266+
267+
@patch("datacommons_client.endpoints.node.flatten_ancestry")
268+
@patch("datacommons_client.endpoints.node.build_ancestry_map")
269+
def test_fetch_entity_ancestry_flat(mock_build_map, mock_flatten):
270+
"""Test fetch_entity_ancestry with flat structure (as_tree=False)."""
271+
mock_build_map.return_value = (
272+
"X",
273+
{
274+
"X": [Parent("A", "A name", "Country")],
275+
"A": [],
276+
},
277+
)
278+
mock_flatten.return_value = [{
279+
"dcid": "A",
280+
"name": "A name",
281+
"type": "Country"
282+
}]
283+
284+
endpoint = NodeEndpoint(api=MagicMock())
285+
result = endpoint.fetch_entity_ancestry("X", as_tree=False)
286+
287+
assert result == {"X": [{"dcid": "A", "name": "A name", "type": "Country"}]}
288+
mock_build_map.assert_called_once()
289+
mock_flatten.assert_called_once()
290+
291+
292+
@patch("datacommons_client.endpoints.node.build_ancestry_tree")
293+
@patch("datacommons_client.endpoints.node.build_ancestry_map")
294+
def test_fetch_entity_ancestry_tree(mock_build_map, mock_build_tree):
295+
"""Test fetch_entity_ancestry with tree structure (as_tree=True)."""
296+
mock_build_map.return_value = (
297+
"Y",
298+
{
299+
"Y": [Parent("Z", "Z name", "Region")],
300+
"Z": [],
301+
},
302+
)
303+
mock_build_tree.return_value = {
304+
"dcid":
305+
"Y",
306+
"name":
307+
None,
308+
"type":
309+
None,
310+
"parents": [{
311+
"dcid": "Z",
312+
"name": "Z name",
313+
"type": "Region",
314+
"parents": []
315+
}],
316+
}
317+
318+
endpoint = NodeEndpoint(api=MagicMock())
319+
result = endpoint.fetch_entity_ancestry("Y", as_tree=True)
320+
321+
assert "Y" in result
322+
assert result["Y"]["dcid"] == "Y"
323+
assert result["Y"]["parents"][0]["dcid"] == "Z"
324+
mock_build_map.assert_called_once()
325+
mock_build_tree.assert_called_once_with("Y", mock_build_map.return_value[1])

0 commit comments

Comments
 (0)