def METHOD_NAME(pb_session):
    model_session = ModelSession(session=pb_session, factory=mock.Mock())
    assert model_session.get_halos(axes="xy") == {"output": (256, 128)}
    assert model_session.get_halos(axes="yx") == {"output": (128, 256)}
    assert model_session.get_halos(axes="yxc") == {"output": (128, 256, 1)}