ncov-robot/robot/chrome.py

from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import logging

from conf import path
from . import url, loc


class Robot:
    def __init__(self):
        service = Service(path.chromedriver, log_path=path.chromedriver_log)
        self.driver = webdriver.Chrome(service=service)
        self.driver.implicitly_wait(10)
        self.wait = WebDriverWait(self.driver, 10)
        
    def wait_to_click(self, loc):
        return self.wait.until(EC.element_to_be_clickable(loc))
    
    def login(self, username, password):
        logging.info('login...')
        self.driver.get(url.logout)
        self.driver.get(url.login)
        self.driver.switch_to.frame('loginIframe')
        self.wait_to_click(loc.password_login_tab).click()
        self.wait_to_click(loc.username).send_keys(username)
        self.wait_to_click(loc.password).send_keys(password + Keys.RETURN)
        self.wait.until(EC.url_to_be(url.uc_index))
        logging.info('login successful')
    
    def punch(self, username, password, at_school, risky, geolocation):
        self.login(username, password)
        self.driver.execute_cdp_cmd('Emulation.setGeolocationOverride', geolocation)
        self.driver.get(url.ncov)
        school_loc = loc.is_at_school if at_school else loc.not_at_school
        risk_loc = loc.risky if risky else loc.no_risk
        self.wait_to_click(school_loc).click()
        self.wait_to_click(risk_loc).click()
        self.wait_to_click(loc.geolocation).click()
        self.wait.until(EC.invisibility_of_element_located((loc.loading)))
        btn = self.wait_to_click(loc.submit_button)
        btn.click()
        if not btn.get_attribute('class'):
            self.wait_to_click(loc.submit_confirm_button).click()
            self.wait.until(EC.invisibility_of_element_located(loc.loading))
        result = self.wait.until(EC.visibility_of_element_located(loc.result)).text
        return result
    
    def quit(self):
        self.driver.quit()
    
    def test_geolocation(self, geolocation):
        self.driver.get(url.geolocation)
        self.driver.execute_cdp_cmd('Emulation.setGeolocationOverride', geolocation)
        self.driver.switch_to.frame('iframeResult')
        self.wait_to_click((loc.By.TAG_NAME, 'button')).click()
        from time import sleep
        sleep(5)
        result = self.driver.find_element(loc.By.ID, 'demo').text
        return result