交通信号灯

'''
交通信号灯显示颜色：红灯后是绿灯，绿灯后是黄灯，黄灯后是红灯
'''
def update_light(current):
    current_list = ["red","green","yellow"]
    n = current_list.index(current)
    if n ==(len(current_list)-1):
        return current_list[0]
    else:
        return current_list[n+1]

def test_light_colour():
    assert update_light("red")=="green"
    assert update_light("green")=="yellow"
    assert update_light("yellow") == "red"