Python

셀레니움 크롬드라이버로 local storage제어하기

순돌아범 2022. 11. 28. 17:34

RPA 및 크롤링 도구로 selenium을 많이 사용하는데, 값을 재사용하기 위해 localstorage 제어하여 활용하였습니다.

​

셀레니움에서 localstorage를 제어하는 부분을 살펴보겠습니다.

​

기본 동작은 아주 간단한데 자바스크립트를 실행시켜 localstorage 키에 값을 넣는 방식입니다.

 

driver.execute_script("window.localStorage.setItem(arguments[0], arguments[1]);", key, value)

 

​

크롬 드라이버 제어 클래스

 

제가 사용하는 크롬 드라이버 클래스 소스의 일부를 발췌하였습니다.

​

chromedrivcer.py

import os
import platform
from os import path
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.common.exceptions import WebDriverException as WDE


class ChromeDriver:
    """
    크롬 드라이버 세팅
    """
    def __init__(self):
        self.__chrome_path = None
        self.driver = None
        self.__option_width = None
        self.__option_height = None
        self.__set_chrome_path()

    def __set_chrome_path(self) -> None:
        """ OS에 따라 크롬 드라이버 위치 지정
            윈도우는 사용하는 크롬브라우저 버전에 맞는 드라이버 버전을 사용하고
            서버에서는 crhome 97버전이 설치되어 있으므로 크롬드라이버도 97로 사용
        """
        if platform.system() == 'Windows':
            self.__chrome_path = path.dirname(path.abspath(__file__)) + "/support/chromedriver.exe"
        elif platform.system() == 'Darwin':
            self.__chrome_path = path.dirname(path.abspath(__file__)) + "/support/chromedriver"
        else:
            self.__chrome_path = path.dirname(path.abspath(__file__)) + "/support/chromedriver_linux97"


    def set_size_option(self, width: int, height: int) -> bool:
        """ driver_ready 전 창 크기 지정
        """
        if not isinstance(width, int) or not isinstance(height, int):
            return False

        self.__option_width = width
        self.__option_height = height
        return True

    def driver_ready(self) -> None:
        """ 크롬 드라이버 초기화"""
        options = webdriver.ChromeOptions()
        options.add_argument('--disable-extensions')
        options.add_argument('--headless')
        options.add_argument('--disable-gpu')
        options.add_argument('--no-sandbox')
        options.add_experimental_option('excludeSwitches', ['enable-logging'])
        if self.__option_width and self.__option_height:
            options.add_argument("--window-size=" + str(self.__option_width) + "," + str(self.__option_height))
        s = Service(self.__chrome_path)
        self.driver = webdriver.Chrome(service=s, options=options)

        self.driver.implicitly_wait(2)

    def set_maximize_window(self) -> None:
        """ 창 크기를 최대로 설정 """
        self.driver.maximize_window()

    def load_page(self, url: str) -> bool:
        """ 지정된 URL을 접속 """
        if not isinstance(url, str):
            return False
        if url.strip("") == "":
            return False
        if url.lower()[0:4] != "http":
            return False

        try:
            self.driver.get(url)
            self.driver.implicitly_wait(2)

        except WDE as ex:
            print("chromedriver > load_page > WDE error : " + url + ", ex : "
                                    + str(ex))
            return False

        return True

    def __get_auto_size(self) -> list:
        """ 화면 내용에 표현되는 사이즈 측정 """
        totalWidth = self.driver.execute_script("return document.body.offsetWidth")
        totalHeight = self.driver.execute_script('return document.body.parentNode.scrollHeight') + 20
        if isinstance(totalWidth, int) and isinstance(totalWidth, int):
            return [totalWidth, totalHeight]
        else:
            return [0, 0]

    def auto_resize(self) -> None:
        """ 화면 내용에 표현되는 사이즈 측정 후 자동으로 크기 지정 """
        size = self.__get_auto_size()
        if size[0] > 20 and size[1] > 20:
            self.resize(size[0], size[1])

    def get_localstorage(self, key: str) -> bool:
        """ localstorage 에서 해당 키값의 정보를 조회 """
        if not isinstance(key, str):
            return False
        if key == '':
            return False

        return self.driver.execute_script("return window.localStorage.getItem(arguments[0]);", key)

    def set_localstorage(self, key: str, value: str) -> bool:
        """ localstorage에 정보를 저장 """
        if not isinstance(key, str):
            return False
        if key == '':
            return False

        if not isinstance(value, str):
            return False

        self.driver.execute_script("window.localStorage.setItem(arguments[0], arguments[1]);", key, value)
        result_value = self.get_localstorage(key)

        if value == result_value:
            return True
        else:
            return False

    def remove_localstorage(self, key: str) -> bool:
        """ localstorage에서 삭제"""

        if not isinstance(key, str):
            return False
        if key == '':
            return False

        self.driver.execute_script("window.localStorage.removeItem(arguments[0]);", key)
        result_value = self.get_localstorage(key)
        if result_value is None:
            return True
        else:
            return False

    def close(self):
        """ 드라이버 닫기 """
        if self.driver is not None:
            self.driver.quit()

 

 

테스트 코드

 

잘 동작하는 지 테스트 해봐야 하므로 테스트 케이스를 작성하였습니다.

​

test_chromedriver.py

import sys
from os import path
import unittest
import time
import os

from chromedriver import ChromeDriver


class MyTestCase(unittest.TestCase):
    """ chromedriver 테스트 클래스 """

    @classmethod
    def setUpClass(cls):
        """ 최초 1회 세팅 """
        cls.driver = ChromeDriver()

    @classmethod
    def tearDownClass(cls):
        """ 최종 세팅 """
        cls.driver.close()


    def test_driver_ready(self):
        # given
        # when
        self.driver.driver_ready()

        # then
        self.assertIsNotNone(self.driver.driver)


    def test_load_page(self):
        """ 페이지 로딩 테스트 """
        # given
        if self.driver.driver is None:
            self.driver.driver_ready()

        # when
        self.driver.load_page("https://google.com")
        html = self.driver.driver.page_source

        # then
        self.assertIsNotNone(html)

    def test_load_page_fail(self):
        """ 페이지 로딩 테스트 실패 케이스 """
        # given
        if self.driver.driver is None:
            self.driver.driver_ready()

        # when

        # then
        self.assertFalse(self.driver.load_page(""))
        self.assertFalse(self.driver.load_page("google.com"))
        self.assertFalse(self.driver.load_page("1234"))
        # self.assertFalse(self.driver.load_page("http://gooooooooooooooooooooooooooooooogle.com"))


    def test_get_localstorage(self):
        """ localstorage에 값 조회 테스트 """
        # given
        self.driver.load_page("https://google.com")
        self.driver.remove_localstorage('uid')

        # when
        result = self.driver.get_localstorage('uid')

        # then
        self.assertIsNone(result)

    def test_get_localstorage_fail(self):
        """ localstorage에 값 조회 테스트 실패 케이스 """
        # given
        self.driver.load_page("https://google.com")

        # when
        result1 = self.driver.get_localstorage(123)
        result2 = self.driver.get_localstorage('')
        result3 = self.driver.get_localstorage(list())
        result4 = self.driver.get_localstorage(dict())
        result5 = self.driver.get_localstorage(False)

        # then
        self.assertFalse(result1)
        self.assertFalse(result2)
        self.assertFalse(result3)
        self.assertFalse(result4)
        self.assertFalse(result5)

    def test_set_localstorage(self):
        """ localstorage에 값 저장 테스트 """
        # given
        uid = '111111-2222-3333-4444-chromedriver'
        self.driver.load_page("https://google.com")

        # when
        result = self.driver.set_localstorage('uid', uid)
        driver_uid = self.driver.get_localstorage('uid')
        result2 = self.driver.set_localstorage('test', '')
        driver_test = self.driver.get_localstorage('test')

        # then
        self.assertTrue(result)
        self.assertEqual(driver_uid, uid)

        self.assertTrue(result2)
        self.assertEqual(driver_test,'')

    def test_set_localstorage_fail(self):
        """ localstorage에 값 조회 테스트 """
        # given
        self.driver.load_page("https://google.com")

        # when
        result1 = self.driver.set_localstorage('uid', None)
        result2 = self.driver.set_localstorage('uid', 123)
        result3 = self.driver.set_localstorage(123, '123')
        result4 = self.driver.set_localstorage(list(), list())
        result5 = self.driver.set_localstorage('', 'fdasfda')
        result6 = self.driver.set_localstorage(dict(), '')

        # then
        self.assertFalse(result1)
        self.assertFalse(result2)
        self.assertFalse(result3)
        self.assertFalse(result4)
        self.assertFalse(result5)
        self.assertFalse(result6)

    def test_remove_localstorage(self):
        """ localstorage에서 삭제 테스트 """
        # given
        self.driver.load_page("https://google.com")
        uid = '111111-2222-3333-4444-chromedriver'
        self.driver.set_localstorage('uid', uid)

        # when
        result = self.driver.remove_localstorage('uid')
        driver_uid = self.driver.get_localstorage('uid')

        # then
        self.assertTrue(result)
        self.assertIsNone(driver_uid)

    def test_remove_localstorage_fail(self):
        """ localstorage에 값 삭제 실패 케이스 """
        # given
        self.driver.load_page("https://google.com")

        # when
        result1 = self.driver.remove_localstorage(123)
        result2 = self.driver.remove_localstorage('')
        result3 = self.driver.remove_localstorage(list())
        result4 = self.driver.remove_localstorage(dict())
        result5 = self.driver.remove_localstorage(False)

        # then
        self.assertFalse(result1)
        self.assertFalse(result2)
        self.assertFalse(result3)
        self.assertFalse(result4)
        self.assertFalse(result5)


if __name__ == '__main__':
    unittest.main(warnings='ignore')

 

 

실행 테스트

 

pytest 로 테스트 케이스들을 테스트 해보니 잘 동작하네요.

​

​

참고 자료