Seleniumで画面遷移した後の画面で要素が取得できない問題の解消

最近PythonでSeleniumを使い始めたんですが、
ちょっとハマったことがあったのでメモ程度に残しておきます。

元のコード

login.py

from selenium import webdriver

driver = webdriver.Chrome()

driver.get("サイトのログインページ")

elem_search_word = driver.find_element_by_css_selector("IDのフォーム")
elem_search_word.send_keys(id) # IDを入力する

elem_search_word = driver.find_element_by_css_selector("PASSWORDのフォーム")
elem_search_word.send_keys(password) # パスワードを入力する
elem_search_btn = driver.find_element_by_css_selector("ログインボタン")
elem_search_btn.click()

elem_text = driver.find_element_by_css_selector("ログイン後のページの要素")
...

こんな感じで、
・サイトのログインページに飛んで、
・フォームにIDとPASSを自力で入力して
・ログインボタンを押してログイン後のページに飛んで
・ログイン後画面の要素を取得する
というプログラムを書いていました。

selenium.common.exceptions.NoSuchElementException: Message: no such element: Unable to locate element: {"method":"css selector","selector":"ログイン後ページの要素"}

ところが、このプログラムを実行すると上記のエラーが発生。
ちゃんとログイン後ページのCSSセレクタは合ってるのに見つからないと怒られてしまう。

原因を調査する

試しにCSSセレクタのところをXPathとかにしてみてもやはりダメだったので別の手段を取ることに。

driver.current_url

というメソッドを使うことでその時点でのURLを取得できるということを発見。

login.py

from selenium import webdriver

driver = webdriver.Chrome()

driver.get("サイトのログインページ")

print(driver.current_url) // URLを確認する

elem_search_word = driver.find_element_by_css_selector("IDのフォーム")
elem_search_word.send_keys(id) # IDを入力する

elem_search_word = driver.find_element_by_css_selector("PASSWORDのフォーム")
elem_search_word.send_keys(password) # パスワードを入力する
elem_search_btn = driver.find_element_by_css_selector("ログインボタン")
elem_search_btn.click()

print(driver.current_url) // URLを確認する

elem_text = driver.find_element_by_css_selector("ログイン後のページの要素")
...

ということで、こんな感じでプログラムを書き直してちゃんとURLが合っているかを確認してみた。
そうすると、
1つ目は期待値通りログインページだったものの
2つ目も同じくログインページのままになっていた。(本当はログイン後のページであってほしかった)

この辺りで、画面が切り替わる前にログイン後ページ向けの処理が行われてしまっていることをようやく理解できた。

修正後のコード

login.py

from selenium import webdriver
from time import sleep　// 新しくインポート

driver = webdriver.Chrome()

driver.get("サイトのログインページ")

elem_search_word = driver.find_element_by_css_selector("IDのフォーム")
elem_search_word.send_keys(id) # IDを入力する

elem_search_word = driver.find_element_by_css_selector("PASSWORDのフォーム")
elem_search_word.send_keys(password) # パスワードを入力する
elem_search_btn = driver.find_element_by_css_selector("ログインボタン")
elem_search_btn.click()

sleep(5) // 読み込みを待つために５秒間処理を止める

elem_text = driver.find_element_by_css_selector("ログイン後のページの要素")
...

画面が切り替わってから要素の読み込みを行えば問題ないだろうという仮説が立ったので、
sleep()を使うことで読み込みのための空き時間を作った。

こうすることで無事ログイン後ページを読み込んだ上で要素取得の処理が走り、
意図した動作をするようになった。