自动化测试登录后的功能

#保存cookies和local storage - 
使用上下文管理器处理文件
def save_cookies(self):
# cookies.pkl存储Cookies信息
    cookies = self.driver.get_cookies()
    with open("cookies.pkl", "wb") as f:
        pickle.dump(cookies, f)
        script = "return JSON.stringify(window.localStorage);"
        # localstorage.pkl存储本地存储库信息
        localStorage = self.excute_js_script(script)
    with open("localstorage.pkl", "wb") as f:
        pickle.dump(localStorage, f)

    # 恢复Cookies和Local Storage
    def load_cookies(self):
         try:
            with open("cookies.pkl", "rb") as f:
                cookies = pickle.load(f)
                for cookie in cookies:
                    self.driver.add_cookie(cookie) 
            with open("localstorage.pkl", "rb") as f:
                localstorage_str = pickle.load(f)
                script = f"""
                (function(storage){{
                    try {{
                        const entries = JSON.parse(storage);
                        for (const [key, value] of Object.entries(entries)) {{
                            window.localStorage.setItem(key, value);
                        }}
                    }} catch(e) {{
                        console.error('LocalStorage恢复失败:', e);
                    }}
                }})({repr(localstorage_str)});
                """
                self.excute_js_script(script)
            # 刷新页面使状态生效
            self.driver.refresh()
         except Exception as e:
            print(f"状态恢复失败: {str(e)}")
            #获取失败，获得截屏页面
            self.driver.save_screenshot("state_restore_failed.png")

    # 判断元素是否存在，若存在赋给某个变量
    def judgeElement(self,way,tag):
        # 添加显式等待确保页面元素加载完成
        wait = WebDriverWait(self.driver, 1)
        return wait.until(EC.element_to_be_clickable((way, tag)))

    #执行JavaScipt语句
    def excute_js_script(self,script):
        self.driver.execute_script(script)

   def setUp(self):
        #这里使用chrome测试，需要指定chromedriver.exe路径
        driverPath = "C:\\Lib\\chromedriver.exe"
        service = Service(executable_path=driverPath)
        options = webdriver.ChromeOptions()
        # 添加用户数据目录以保持会话
        options.add_argument("user-data-dir=C:\\Lib\\ChromeProfile")
        self.driver = webdriver.Chrome(service=service, options=options)
        self.driver.get("http://127.0.0.1:8000")
        #检查cookies.pkl是否存在
        if not os.path.isfile("cookies.pkl"):
            try:
                # 处理可能的登录表单
                element_username = self.judgeElement(By.ID, "id_username")
                element_password = self.judgeElement(By.ID, "id_password")             
                # 尝试多种提交按钮定位方式
                try:
                    element_submit = self.judgeElement(By.NAME, "form-signin")
                except:
                    try:
                        element_submit = self.judgeElement(By.CSS_SELECTOR, "button[type='submit']")
                    except:
                        element_submit = self.judgeElement(By.XPATH, "//button[contains(.,'登录') or contains(.,'Sign in')]")                        
                # 执行登录
                element_username.clear()
                element_username.send_keys("cindy")
                element_password.clear()
                element_password.send_keys("123456")
                element_submit.click()
            except Exception as e:
                print(f"登录可能已完成或失败: {str(e)}")
                self.driver.save_screenshot("login_error.png")
            # 验证登录是否成功
            try:
                wait.until(EC.title_contains("电子商务系统"))
                if self.driver.title != "电子商务系统":
                    self.fail("标题不正确，实际为: " + self.driver.title)
                #登录成功后存储cookies信息
                self.save_cookies()
            except:
                self.driver.save_screenshot("title_verification_failed.png")
                raise

def test_search(self):
        #获取Cookies信息
        self.load_cookies()
        self.driver.get("http://127.0.0.1:8000/goods_view/")
        try:
            # 执行搜索
            search_input = self.judgeElement(By.NAME,"good")
            search_input.clear()
            search_input.send_keys("茶")
            submit_button =self.judgeElement(By.XPATH,"//*[@id='navbar']/form/button")
            submit_button.click()
            elements = self.driver.find_elements(By.LINK_TEXT, "放入")
            self.assertTrue(len(elements)> 0, "没有查询到")
        except Exception as e:
            print(f"搜索失败: {str(e)}")
            self.driver.save_screenshot("search_failed.png")
            raise

def test_put_into_char(self):
    self.load_cookies()
    self.driver.get("http://127.0.0.1:8000/goods_view/")
    try:
          input_link = self.judgeElement(By.LINK_TEXT,"放入")
          input_link.click()
          elements = self.driver.find_elements(By.XPATH,'//font[@color="#FF0000" and text()="1"]')
          self.assertIsNotNone(elements)
    except Exception as e:
          print(f"放入失败: {str(e)}")
          self.driver.save_screenshot("put_into_char_failed.png")
          raise

def tearDown(self):
  self.driver.quit()
if __name__ == "__main__": 
  unittest.main()