From 218082fed6529e8a30a40bee12efddbb76cc20d2 Mon Sep 17 00:00:00 2001
From: Jerry Yan <792602257@qq.com>
Date: Mon, 6 Jun 2022 15:07:07 +0800
Subject: [PATCH] =?UTF-8?q?op=5Ftab=E5=BF=AB=E6=8D=B7=E6=96=B9=E6=B3=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 browser/IDriver.py | 20 ++++++++++++++++++++
 browser/chrome.py  | 20 +++++++++++---------
 browser/edge.py    | 20 +++++++++++---------
 3 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/browser/IDriver.py b/browser/IDriver.py
index 4eff625..cb8b353 100644
--- a/browser/IDriver.py
+++ b/browser/IDriver.py
@@ -1,12 +1,32 @@
+import contextlib
+
+
 class IDriver():
     def new_tab(self) -> str:
         ...
 
+    def get_current_tab(self) -> str:
+        ...
+
     def change_tab(self, tab_handler: str):
         ...
 
     def open_url(self, url: str, tab_handler: str = ""):
         ...
 
+    @contextlib.contextmanager
+    def op_tab(self, tab_handler: str):
+        cur_handle = self.get_current_tab()
+        if tab_handler == "":
+            tab_handler = cur_handle
+        try:
+            self.change_tab(tab_handler)
+            yield self
+        finally:
+            self.change_tab(cur_handle)
+
+    def refresh(self, tab_handler: str = ""):
+        ...
+
     def screenshot(self, tab_handler: str = "") -> str:
         ...
diff --git a/browser/chrome.py b/browser/chrome.py
index b99ada5..939a3c9 100644
--- a/browser/chrome.py
+++ b/browser/chrome.py
@@ -45,15 +45,17 @@ class ChromeDriver(IDriver):
             return
         self.browser.switch_to.window(tab_handler)
 
+    def get_current_tab(self) -> str:
+        return self.browser.current_window_handle
+
     def open_url(self, url: str, tab_handler: str = ""):
-        if tab_handler != "":
-            self.change_tab(tab_handler)
-        self.browser.get(url)
+        with self.op_tab(tab_handler):
+            self.browser.get(url)
+
+    def refresh(self, tab_handler: str = ""):
+        with self.op_tab(tab_handler):
+            self.browser.refresh()
 
     def screenshot(self, tab_handler: str = "") -> str:
-        cur_handle = self.browser.current_window_handle
-        if tab_handler != "":
-            self.change_tab(tab_handler)
-        b64 = self.browser.get_screenshot_as_base64()
-        self.change_tab(cur_handle)
-        return b64
+        with self.op_tab(tab_handler):
+            return self.browser.get_screenshot_as_base64()
diff --git a/browser/edge.py b/browser/edge.py
index 28648fc..4bb2e5a 100644
--- a/browser/edge.py
+++ b/browser/edge.py
@@ -38,6 +38,9 @@ class EdgeDriver(IDriver):
                 return _handle
         return ""
 
+    def get_current_tab(self) -> str:
+        return self.browser.current_window_handle
+
     def change_tab(self, tab_handler: str):
         if tab_handler not in self.browser.window_handles:
             return
@@ -46,14 +49,13 @@ class EdgeDriver(IDriver):
         self.browser.switch_to.window(tab_handler)
 
     def open_url(self, url: str, tab_handler: str = ""):
-        if tab_handler != "":
-            self.change_tab(tab_handler)
-        self.browser.get(url)
+        with self.op_tab(tab_handler):
+            self.browser.get(url)
+
+    def refresh(self, tab_handler: str = ""):
+        with self.op_tab(tab_handler):
+            self.browser.refresh()
 
     def screenshot(self, tab_handler: str = "") -> str:
-        cur_handle = self.browser.current_window_handle
-        if tab_handler != "":
-            self.change_tab(tab_handler)
-        b64 = self.browser.get_screenshot_as_base64()
-        self.change_tab(cur_handle)
-        return b64
+        with self.op_tab(tab_handler):
+            return self.browser.get_screenshot_as_base64()