def test_ElectrodeArray_add_electrode():
    earray = ElectrodeArray([])
    npt.assert_equal(earray.n_electrodes, 0)

    with pytest.raises(TypeError):
        earray.add_electrode('A01', ElectrodeArray([]))

    # Add an electrode:
    key0 = 'A04'
    earray.add_electrode(key0, PointSource(0, 1, 2))
    npt.assert_equal(earray.n_electrodes, 1)
    # Both numeric and string index should work:
    for key in [key0, 0]:
        npt.assert_equal(isinstance(earray[key], PointSource), True)
        npt.assert_almost_equal(earray[key].x, 0)
        npt.assert_almost_equal(earray[key].y, 1)
        npt.assert_almost_equal(earray[key].z, 2)
    with pytest.raises(ValueError):
        # Can't add the same electrode twice:
        earray.add_electrode(key0, PointSource(0, 1, 2))

    # Add another electrode:
    key1 = 'A01'
    earray.add_electrode(key1, DiskElectrode(4, 5, 6, 7))
    npt.assert_equal(earray.n_electrodes, 2)
    # Both numeric and string index should work:
    for key in [key1, 1]:
        npt.assert_equal(isinstance(earray[key], DiskElectrode), True)
        npt.assert_almost_equal(earray[key].x, 4)
        npt.assert_almost_equal(earray[key].y, 5)
        npt.assert_almost_equal(earray[key].z, 6)
        npt.assert_almost_equal(earray[key].r, 7)

    # We can also get a list of electrodes:
    for keys in [[key0, key1], [0, key1], [key0, 1], [0, 1]]:
        selected = earray[keys]
        npt.assert_equal(isinstance(selected, list), True)
        npt.assert_equal(isinstance(selected[0], PointSource), True)
        npt.assert_equal(isinstance(selected[1], DiskElectrode), True)
def test_ElectrodeArray_remove_electrode():
    earray1 = ElectrodeArray([])
    earray2 = ElectrodeArray([])
    npt.assert_equal(earray1.n_electrodes, 0)

    # Can't remove electrodes from empty electrodeArray
    with pytest.raises(ValueError):
        earray1.remove_electrode(None)
    with pytest.raises(ValueError):
        earray1.remove_electrode("foo")

    key = [0] * 4
    key[0] = 'D03'
    key[1] = 'A02'
    key[2] = 'F10'
    key[3] = 'E12'

    earray1.add_electrode(key[0], PointSource(0, 1, 2))
    earray1.add_electrode(key[1], PointSource(3, 4, 5))
    earray1.add_electrode(key[2], PointSource(6, 7, 8))
    earray1.add_electrode(key[3], PointSource(9, 10, 11))
    npt.assert_equal(earray1.n_electrodes, 4)

    earray2.add_electrode(key[0], PointSource(0, 1, 2))
    earray2.add_electrode(key[1], PointSource(3, 4, 5))
    earray2.add_electrode(key[2], PointSource(6, 7, 8))
    earray2.add_electrode(key[3], PointSource(9, 10, 11))
    npt.assert_equal(earray2.n_electrodes, 4)

    # Remove one electrode key[1] from the electrodeArray
    earray1.remove_electrode(key[0])
    npt.assert_equal(earray1.n_electrodes, 3)
    # Can't remove an electrode that has been removed
    with pytest.raises(ValueError):
        earray1.remove_electrode(key[0])

    # List keeps order:
    npt.assert_equal(earray1[0], earray1[key[1]])
    npt.assert_equal(earray1[1], earray1[key[2]])
    npt.assert_equal(earray1[2], earray1[key[3]])

    # Other electrodes stay the same
    for k in [key[1], key[2], key[3]]:
        npt.assert_equal(earray1[k].x, earray2[k].x)
        npt.assert_equal(earray1[k].y, earray2[k].y)
        npt.assert_equal(earray1[k].z, earray2[k].z)

    # Remove two more electrodes from the electrodeArray
    # List keeps order
    earray1.remove_electrode(key[1])
    earray1.remove_electrode(key[2])
    npt.assert_equal(earray1.n_electrodes, 1)
    npt.assert_equal(earray1[0], earray1[key[3]])

    # The last electrode stays the same
    for key in [key[3]]:
        npt.assert_equal(earray1[key].x, earray2[key].x)
        npt.assert_equal(earray1[key].y, earray2[key].y)
        npt.assert_equal(earray1[key].z, earray2[key].z)