From 7579c10cb71571f865ba0f85ed897043ee767e63 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Mon, 17 Oct 2022 11:47:20 +0800
Subject: [PATCH] feat: add retry when requesting ocr service

---
 hrp/pkg/uixt/ocr_vedem.go | 30 +++++++++++++++++++++---------
 httprunner/utils.py       | 11 ++++++++---
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/hrp/pkg/uixt/ocr_vedem.go b/hrp/pkg/uixt/ocr_vedem.go
index 82d142dc..91c73e2f 100644
--- a/hrp/pkg/uixt/ocr_vedem.go
+++ b/hrp/pkg/uixt/ocr_vedem.go
@@ -45,7 +45,7 @@ func (s *veDEMOCRService) getOCRResult(imageBuf []byte) ([]OCRResult, error) {
 	if err != nil {
 		return nil, fmt.Errorf("create form file error: %v", err)
 	}
-	_, err = formWriter.Write(imageBuf)
+	size, err := formWriter.Write(imageBuf)
 	if err != nil {
 		return nil, fmt.Errorf("write form error: %v", err)
 	}
@@ -55,8 +55,9 @@ func (s *veDEMOCRService) getOCRResult(imageBuf []byte) ([]OCRResult, error) {
 		return nil, fmt.Errorf("close body writer error: %v", err)
 	}
 
-	if env.VEDEM_OCR_URL == "" {
-		log.Error().Msg("VEDEM_OCR_URL env missed for OCR service")
+	if env.VEDEM_OCR_URL == "" || env.VEDEM_OCR_AK == "" || env.VEDEM_OCR_SK == "" {
+		log.Error().Msg(
+			"missed env missed for veDEM OCR service: VEDEM_OCR_URL/VEDEM_OCR_AK/VEDEM_OCR_SK")
 		os.Exit(1)
 	}
 
@@ -69,14 +70,25 @@ func (s *veDEMOCRService) getOCRResult(imageBuf []byte) ([]OCRResult, error) {
 
 	req.Header.Add("Agw-Auth", token)
 	req.Header.Add("Content-Type", bodyWriter.FormDataContentType())
-	resp, err := client.Do(req)
-	if err != nil {
-		var logID string
-		if resp != nil {
-			logID = getLogID(resp.Header)
+
+	var resp *http.Response
+	// retry 3 times
+	for i := 1; i <= 3; i++ {
+		resp, err = client.Do(req)
+		if err == nil {
+			break
 		}
-		return nil, fmt.Errorf("http reqeust veDEM OCR server error: %v, logID: %s", err, logID)
+		logID := getLogID(resp.Header)
+		log.Error().Err(err).
+			Str("logID", logID).
+			Int("imageBufSize", size).
+			Msgf("request OCR service failed, retry %d", i)
+		time.Sleep(1 * time.Second)
 	}
+	if resp == nil {
+		return nil, fmt.Errorf("veDEM OCR service is not available")
+	}
+
 	defer resp.Body.Close()
 
 	results, err := ioutil.ReadAll(resp.Body)
diff --git a/httprunner/utils.py b/httprunner/utils.py
index b95d3ecb..cbcc8f8e 100644
--- a/httprunner/utils.py
+++ b/httprunner/utils.py
@@ -243,7 +243,8 @@ def sort_dict_by_custom_order(raw_dict: Dict, custom_order: List):
 
 
 class ExtendJSONEncoder(json.JSONEncoder):
-    """especially used to safely dump json data with python object, such as MultipartEncoder"""
+    """especially used to safely dump json data with python object,
+    such as MultipartEncoder"""
 
     def default(self, obj):
         try:
@@ -275,7 +276,8 @@ def is_support_multiprocessing() -> bool:
         Queue()
         return True
     except (ImportError, OSError):
-        # system that does not support semaphores(dependency of multiprocessing), like Android termux
+        # system that does not support semaphores
+        # (dependency of multiprocessing), like Android termux
         return False
 
 
@@ -320,7 +322,10 @@ def gen_cartesian_product(*args: List[Dict]) -> List[Dict]:
     return product_list
 
 
-LOGGER_FORMAT = "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level}</level> | <level>{message}</level>"
+LOGGER_FORMAT = (
+    "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green>"
+    + " | <level>{level}</level> | <level>{message}</level>"
+)
 
 
 def init_logger(level: str):