* Directly call async_write_ha_state pt2 * Directly call async_write_ha_state pt2 * Fix mock * Address comments