146 lines
5.1 KiB
Python
146 lines
5.1 KiB
Python
"""
|
|
Tests for NodeIO and LazyProxy.
|
|
"""
|
|
import pytest
|
|
from unittest.mock import MagicMock
|
|
from pytrace.core.io import NodeIO
|
|
from pytrace.core.proxy import LazyProxy
|
|
from pytrace.core.data import TraceData
|
|
|
|
@pytest.fixture
|
|
def node_io():
|
|
return NodeIO()
|
|
|
|
@pytest.fixture
|
|
def mock_node():
|
|
return MagicMock(uid="mock_uid", name="MockNode", params={})
|
|
|
|
def test_node_io_init(node_io):
|
|
assert node_io._inputs is None
|
|
assert node_io._params is None
|
|
assert node_io._outputs is None
|
|
assert node_io._node_ref is None
|
|
|
|
def test_node_io_bind_and_clear(node_io, mock_node):
|
|
inputs = {"test_in": TraceData(value="data")}
|
|
params = {"test_param": 123}
|
|
node_io.bind(inputs, params, mock_node)
|
|
|
|
assert node_io._inputs == inputs
|
|
assert node_io._params == params
|
|
assert node_io._node_ref == mock_node
|
|
assert node_io._outputs == {}
|
|
|
|
node_io.clear()
|
|
assert node_io._inputs is None
|
|
assert node_io._params is None
|
|
assert node_io._outputs is None # Should be cleared
|
|
assert node_io._node_ref is None
|
|
|
|
def test_node_io_get_input_local(node_io, mock_node):
|
|
local_data = TraceData(value="local_value")
|
|
node_io.bind({"my_input": local_data}, {}, mock_node)
|
|
|
|
result = node_io.get_input("my_input")
|
|
assert result == "local_value"
|
|
|
|
def test_node_io_get_input_remote(node_io, mock_node):
|
|
remote_data = TraceData(ref="remote_ref_123", meta={"metadata_key": "meta_value"})
|
|
node_io.bind({"my_remote_input": remote_data}, {}, mock_node)
|
|
|
|
result = node_io.get_input("my_remote_input")
|
|
assert isinstance(result, LazyProxy)
|
|
|
|
# Test accessing metadata through proxy
|
|
assert result.metadata_key == "meta_value"
|
|
|
|
# Test force_load on LazyProxy (mock fetch)
|
|
with pytest.raises(AttributeError): # LazyProxy tries to fetch via capsule.fetch, which raises AttributeError by default
|
|
result.force_load()
|
|
|
|
def test_node_io_get_input_default(node_io, mock_node):
|
|
node_io.bind({}, {}, mock_node)
|
|
result = node_io.get_input("non_existent", "default_val")
|
|
assert result == "default_val"
|
|
|
|
def test_node_io_get_param_from_inputs_override(node_io, mock_node):
|
|
# Param overridden by input
|
|
node_io.bind({"my_param": TraceData(value="input_override")}, {"my_param": "original_param"}, mock_node)
|
|
result = node_io.get_param("my_param")
|
|
assert result == "input_override"
|
|
|
|
def test_node_io_get_param_from_params(node_io, mock_node):
|
|
# Param from static config
|
|
node_io.bind({}, {"my_param": "static_config"}, mock_node)
|
|
result = node_io.get_param("my_param")
|
|
assert result == "static_config"
|
|
|
|
def test_node_io_get_param_default(node_io, mock_node):
|
|
node_io.bind({}, {}, mock_node)
|
|
result = node_io.get_param("non_existent", "default_param_val")
|
|
assert result == "default_param_val"
|
|
|
|
def test_node_io_set_and_get_output(node_io, mock_node):
|
|
node_io.bind({}, {}, mock_node)
|
|
node_io.set_output("result1", "output_data_1")
|
|
node_io.set_output("result2", 123)
|
|
|
|
outputs = node_io.get_collected_outputs()
|
|
assert "result1" in outputs
|
|
assert outputs["result1"].value == "output_data_1"
|
|
assert "result2" in outputs
|
|
assert outputs["result2"].value == 123
|
|
|
|
def test_node_io_empty_outputs(node_io, mock_node):
|
|
node_io.bind({}, {}, mock_node)
|
|
outputs = node_io.get_collected_outputs()
|
|
assert outputs == {}
|
|
|
|
def test_lazy_proxy_local_data_access():
|
|
capsule = TraceData(value={"a": 1, "b": 2}, meta={"key": "val"})
|
|
proxy = LazyProxy(capsule)
|
|
|
|
assert proxy.key == "val" # Access metadata
|
|
|
|
# Direct access to an attribute that exists in materialized data
|
|
# LazyProxy does not materialize for metadata access
|
|
# It materializes if the attribute is not in meta and does not exist in capsule itself
|
|
# For now, if the value is not remote, LazyProxy should not trigger fetch on attribute access
|
|
# However, the current LazyProxy will try to fetch if not in meta. This needs to be refined.
|
|
# For testing, we mock fetch.
|
|
|
|
# Mock the fetch method for TraceData
|
|
capsule.fetch = MagicMock(return_value={"a": 1, "b": 2, "c": 3})
|
|
|
|
assert proxy.a == 1
|
|
assert proxy.b == 2
|
|
assert hasattr(proxy, "c")
|
|
assert proxy.c == 3 # Should have materialized
|
|
capsule.fetch.assert_called_once()
|
|
|
|
|
|
def test_lazy_proxy_force_load():
|
|
capsule = TraceData(ref="remote_ref")
|
|
proxy = LazyProxy(capsule)
|
|
|
|
mock_fetched_data = MagicMock(data="actual data")
|
|
capsule.fetch = MagicMock(return_value=mock_fetched_data)
|
|
|
|
loaded_data = proxy.force_load()
|
|
assert loaded_data == mock_fetched_data
|
|
capsule.fetch.assert_called_once()
|
|
|
|
# Subsequent calls should not fetch again
|
|
loaded_data_2 = proxy.force_load()
|
|
assert loaded_data_2 == mock_fetched_data
|
|
capsule.fetch.assert_called_once()
|
|
|
|
def test_lazy_proxy_repr():
|
|
capsule = TraceData(ref="remote_ref_repr")
|
|
proxy = LazyProxy(capsule)
|
|
assert repr(proxy) == "<LazyProxy for <TraceData remote_ref='remote_ref_repr'>>"
|
|
|
|
capsule_local = TraceData(value="local")
|
|
proxy_local = LazyProxy(capsule_local)
|
|
assert repr(proxy_local) == "<LazyProxy for <TraceData value=str>>"
|