def METHOD_NAME(
    managed_cluster_restarts: ManagedCluster,
) -> Iterator[ManagedCluster]:
    managed_cluster_restarts.wait_for_agent_ok(20)
    try:
        yield managed_cluster_restarts
        managed_cluster_restarts.wait_for_agent_ok(20)
    except Exception:
        managed_cluster_restarts.restart_master()
        managed_cluster_restarts.restart_agent()
        raise