How to use the pywr.nodes.Link function in pywr

def test_virtual_storage_duplicate_route():
    """ Test the VirtualStorage node """

    model = pywr.core.Model()

    inpt = Input(model, "Input", max_flow=20)
    lnk = Link(model, "Link")
    inpt.connect(lnk)
    otpt = Output(model, "Output", max_flow=10, cost=-10.0)
    lnk.connect(otpt)

    vs = pywr.core.VirtualStorage(model, "Licence", [lnk, otpt], factors=[0.5, 1.0], initial_volume=10.0, max_volume=10.0)

    model.setup()

    assert_allclose(vs.volume, [10], atol=1e-7)

    model.step()

    assert_allclose(otpt.flow, [10/1.5], atol=1e-7)
    assert_allclose(vs.volume, [0], atol=1e-7)

    model.step()

Compensation      |
                    |           |
                    v           |
    Catchment -> River 1 -> River 2 ----> MRFA -> Waste
                                    |              ^
                                    |---> MRFB ----|
    """
    model = Model()

    catchment = Input(model, "catchment", max_flow=500, min_flow=500)

    reservoir = Storage(model, "reservoir", max_volume=10000, initial_volume=5000)

    demand = Output(model, "demand", max_flow=50, cost=-100)
    pumping_station = Link(model, "pumping station", max_flow=100, cost=-10)
    river1 = Link(model, "river1")
    river2 = Link(model, "river2")
    compensation = Link(model, "compensation", cost=600)
    mrfA = Link(model, "mrfA", cost=-500, max_flow=50)
    mrfB = Link(model, "mrfB")
    waste = Output(model, "waste")

    catchment.connect(river1)
    river1.connect(river2)
    river2.connect(mrfA)
    river2.connect(mrfB)
    mrfA.connect(waste)
    mrfB.connect(waste)
    river2.connect(pumping_station)
    pumping_station.connect(reservoir)
    reservoir.connect(compensation)
    compensation.connect(river1)

def test_virtual_storage():
    """ Test the VirtualStorage node """

    model = pywr.core.Model()

    inpt = Input(model, "Input", max_flow=20)
    lnk = Link(model, "Link")
    inpt.connect(lnk)
    otpt = Output(model, "Output", max_flow=10, cost=-10.0)
    lnk.connect(otpt)

    vs = pywr.core.VirtualStorage(model, "Licence", [lnk], initial_volume=10.0, max_volume=10.0)

    model.setup()

    assert_allclose(vs.volume, [10], atol=1e-7)

    model.step()

    assert_allclose(otpt.flow, [10], atol=1e-7)
    assert_allclose(vs.volume, [0], atol=1e-7)

    model.step()

|              ^
                                    |---> MRFB ----|
    """
    model = Model()

    catchment = Input(model, "catchment", max_flow=500, min_flow=500)

    reservoir = Storage(model, "reservoir", max_volume=10000, initial_volume=5000)

    demand = Output(model, "demand", max_flow=50, cost=-100)
    pumping_station = Link(model, "pumping station", max_flow=100, cost=-10)
    river1 = Link(model, "river1")
    river2 = Link(model, "river2")
    compensation = Link(model, "compensation", cost=600)
    mrfA = Link(model, "mrfA", cost=-500, max_flow=50)
    mrfB = Link(model, "mrfB")
    waste = Output(model, "waste")

    catchment.connect(river1)
    river1.connect(river2)
    river2.connect(mrfA)
    river2.connect(mrfB)
    mrfA.connect(waste)
    mrfB.connect(waste)
    river2.connect(pumping_station)
    pumping_station.connect(reservoir)
    reservoir.connect(compensation)
    compensation.connect(river1)
    reservoir.connect(demand)

    model.check()
    model.setup()

Catchment -> River 1 -> River 2 ----> MRFA -> Waste
                                    |              ^
                                    |---> MRFB ----|
    """
    model = Model()

    catchment = Input(model, "catchment", max_flow=500, min_flow=500)

    reservoir = Storage(model, "reservoir", max_volume=10000, initial_volume=5000)

    demand = Output(model, "demand", max_flow=50, cost=-100)
    pumping_station = Link(model, "pumping station", max_flow=100, cost=-10)
    river1 = Link(model, "river1")
    river2 = Link(model, "river2")
    compensation = Link(model, "compensation", cost=600)
    mrfA = Link(model, "mrfA", cost=-500, max_flow=50)
    mrfB = Link(model, "mrfB")
    waste = Output(model, "waste")

    catchment.connect(river1)
    river1.connect(river2)
    river2.connect(mrfA)
    river2.connect(mrfB)
    mrfA.connect(waste)
    mrfB.connect(waste)
    river2.connect(pumping_station)
    pumping_station.connect(reservoir)
    reservoir.connect(compensation)
    compensation.connect(river1)
    reservoir.connect(demand)

    model.check()

|           |
                    v           |
    Catchment -> River 1 -> River 2 ----> MRFA -> Waste
                                    |              ^
                                    |---> MRFB ----|
    """
    model = Model()

    catchment = Input(model, "catchment", max_flow=500, min_flow=500)

    reservoir = Storage(model, "reservoir", max_volume=10000, initial_volume=5000)

    demand = Output(model, "demand", max_flow=50, cost=-100)
    pumping_station = Link(model, "pumping station", max_flow=100, cost=-10)
    river1 = Link(model, "river1")
    river2 = Link(model, "river2")
    compensation = Link(model, "compensation", cost=600)
    mrfA = Link(model, "mrfA", cost=-500, max_flow=50)
    mrfB = Link(model, "mrfB")
    waste = Output(model, "waste")

    catchment.connect(river1)
    river1.connect(river2)
    river2.connect(mrfA)
    river2.connect(mrfB)
    mrfA.connect(waste)
    mrfB.connect(waste)
    river2.connect(pumping_station)
    pumping_station.connect(reservoir)
    reservoir.connect(compensation)
    compensation.connect(river1)
    reservoir.connect(demand)

def create_model():
    # create a model
    model = Model(start="2016-01-01", end="2019-12-31", timestep=7)
    
    # create three nodes (an input, a link, and an output)
    A = Input(model, name="A", max_flow=10.0)
    B = Link(model, name="B", cost=10.0)
    C = Output(model, name="C", max_flow=5.0, cost=-20.0)
    
    # connect the nodes together
    A.connect(B)
    B.connect(C)
    
    return model

def __init__(self, *args, **kwargs):
        """Initialise a new Link node

        Parameters
        ----------
        max_flow : float or function (optional)
            A maximum flow constraint on the link, e.g. 5.0
        """
        kwargs['color'] = kwargs.pop('color', '#A0A0A0')  # 45% grey
        super(Link, self).__init__(*args, **kwargs)

def __init__(self, model, name, **kwargs):
        storage_name = "{} (storage)".format(name)
        link_name = "{} (link)".format(name)
        assert(storage_name not in model.nodes)
        assert(link_name not in model.nodes)
        self.storage = Storage(
            model,
            name=storage_name,
            min_volume=0,
            max_volume=0,
            initial_volume=0,
            cost=0,
        )
        self.link = Link(
            model,
            name=link_name
        )

        self.storage.connect(self.link)

        super(BreakLink, self).__init__(model, name, **kwargs)