import time


def wait_for_input_states(medialive, input_id, states):
    """
    Wait for a MediaLive input to be in a specified state.
    """
    current_state = ''
    while current_state not in states:
        time.sleep(5)
        current_state = medialive.describe_input(InputId=input_id)['State']
    return current_state