pytrace/tests/unit/test_io.py

"""
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>>"