From ca28d1bdb3233e27ca2558829d7504b12757c400 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Mon, 1 Jun 2020 17:46:21 +0800
Subject: [PATCH] refactor: StepValidation

---
 .../request_with_functions_test.py            |  2 +-
 httprunner/testcase.py                        | 60 ++++++++++++-------
 2 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/examples/postman_echo/request_methods/request_with_functions_test.py b/examples/postman_echo/request_methods/request_with_functions_test.py
index 520608c1..e084d808 100644
--- a/examples/postman_echo/request_methods/request_with_functions_test.py
+++ b/examples/postman_echo/request_methods/request_with_functions_test.py
@@ -15,13 +15,13 @@ class TestCaseRequestWithFunctions(HttpRunner):
     teststeps = [
         Step("get with params")
         .with_variables(foo1="bar1", foo2="session_bar2", sum_v="${sum_two(1, 2)}")
+        .set_extractor("session_foo2", "body.args.foo2")
         .run_request(
             Request()
             .get("/get")
             .with_params(foo1="$foo1", foo2="$foo2", sum_v="$sum_v")
             .with_headers(**{"User-Agent": "HttpRunner/${get_httprunner_version()}"})
         )
-        .extract("session_foo2", "body.args.foo2")
         .assert_equal("status_code", 200)
         .assert_equal("body.args.foo1", "session_bar1")
         .assert_equal("body.args.sum_v", "3")
diff --git a/httprunner/testcase.py b/httprunner/testcase.py
index 38334879..9e94d3cf 100644
--- a/httprunner/testcase.py
+++ b/httprunner/testcase.py
@@ -1,5 +1,5 @@
 import inspect
-from typing import Text, Any
+from typing import Text, Any, Dict
 
 from httprunner.schema import (
     TConfig,
@@ -126,39 +126,33 @@ class Request(object):
         return RequestWithOptionalArgs(MethodEnum.PATCH, url)
 
 
-class Step(object):
-    def __init__(self, name: Text):
+class StepValidation(object):
+    def __init__(
+        self, name: Text, variables: Dict, extractors: Dict, request: TRequest
+    ):
         self.__name = name
-        self.__variables = {}
-        self.__request = None
-        self.__extract = {}
+        self.__variables = variables
+        self.__extractors = extractors
+        self.__request = request
         self.__validators = []
 
-    def with_variables(self, **variables) -> "Step":
-        self.__variables.update(variables)
-        return self
-
     @property
     def request(self) -> TRequest:
         return self.__request
 
-    def run_request(self, req_obj: RequestWithOptionalArgs) -> "Step":
-        self.__request = req_obj.perform()
-        return self
-
-    def extract(self, var_name: Text, jmes_path: Text) -> "Step":
-        self.__extract[var_name] = jmes_path
-        return self
-
-    def assert_equal(self, jmes_path: Text, expected_value: Any) -> "Step":
+    def assert_equal(self, jmes_path: Text, expected_value: Any) -> "StepValidation":
         self.__validators.append({"eq": [jmes_path, expected_value]})
         return self
 
-    def assert_greater_than(self, jmes_path: Text, expected_value: Any) -> "Step":
+    def assert_greater_than(
+        self, jmes_path: Text, expected_value: Any
+    ) -> "StepValidation":
         self.__validators.append({"gt": [jmes_path, expected_value]})
         return self
 
-    def assert_less_than(self, jmes_path: Text, expected_value: Any) -> "Step":
+    def assert_less_than(
+        self, jmes_path: Text, expected_value: Any
+    ) -> "StepValidation":
         self.__validators.append({"lt": [jmes_path, expected_value]})
         return self
 
@@ -167,6 +161,28 @@ class Step(object):
             name=self.__name,
             variables=self.__variables,
             request=self.__request,
-            extract=self.__extract,
+            extract=self.__extractors,
             validate=self.__validators,
         )
+
+
+class Step(object):
+    def __init__(self, name: Text):
+        self.__name = name
+        self.__variables = {}
+        self.__extractors = {}
+        self.__request = None
+
+    def with_variables(self, **variables) -> "Step":
+        self.__variables.update(variables)
+        return self
+
+    def set_extractor(self, var_name: Text, jmes_path: Text) -> "Step":
+        self.__extractors[var_name] = jmes_path
+        return self
+
+    def run_request(self, req_obj: RequestWithOptionalArgs) -> "StepValidation":
+        self.__request = req_obj.perform()
+        return StepValidation(
+            self.__name, self.__variables, self.__extractors, self.__request
+        )