Flask Integration (#2)

* Added keploy base file

* Added base functions(fetch, test, start, end, put, get) to keploy.py

* Added exception handling and fixed bugs in keploy.py

* Flask Integration Added

* Fixed Denoising

Co-authored-by: Mahesh Gupta <maheshg@mindfiresolutions.com>

keploy/__init__.py

@@ -1,3 +1,5 @@
 from .keploy import Keploy
-from .models import *
-from .mode import setMode
+from .models import Dependency, AppConfig, ServerConfig, FilterConfig, Config, TestCase, TestCaseRequest, TestReq, HttpReq, HttpResp
+from .contrib.flask import KFlask
+from .mode import setMode, getMode
+from .utils import capture_test

keploy/contrib/flask/__init__.py

@@ -0,0 +1,62 @@
+#  BSD 3-Clause License
+
+import copy
+import io
+from typing import Any
+from flask import Flask, request
+from keploy.constants import MODE_OFF
+import keploy as k
+from keploy.contrib.flask.utils import get_request_data
+from keploy.models import HttpResp
+from keploy.utils import capture_test
+from werkzeug import Response
+
+class KFlask(object):
+
+    def __init__(self, keploy:k.Keploy=None, app:Flask=None):
+        self.app = app
+        self.keploy = keploy
+
+        if not app:
+            raise ValueError("Flask app instance not passed, Please initiate flask instance and pass it as keyword argument.")
+
+        if not keploy or k.getMode() == MODE_OFF:
+            return
+
+        app.wsgi_app = KeployMiddleware(keploy, app.wsgi_app)
+
+
+class KeployMiddleware(object):
+    def __init__(self, kep, app) -> None:
+        self.app = app
+        self.keploy = kep
+
+    def __call__(self, environ, start_response) -> Any:
+
+        if not self.keploy:
+            return self.app(environ, start_response)
+
+        req = {}
+        res = {}
+
+        def _start_response(status, response_headers, *args):
+            nonlocal req
+            nonlocal res
+            req = get_request_data(request)
+            res['header'] = {key: [value] for key,value in response_headers}
+            res['status_code'] = int(status.split(' ')[0])
+            return start_response(status, response_headers, *args)
+
+        def _end_response(resp_body):
+            nonlocal res
+            res['body'] = b"".join(resp_body).decode("utf8")
+            return [res['body'].encode('utf-8')]
+
+        resp = _end_response(self.app(environ, _start_response))
+
+        if environ.get("HTTP_KEPLOY_TEST_ID", None):
+            self.keploy.put_resp(environ.get('HTTP_KEPLOY_TEST_ID'), HttpResp(**res))
+        else:
+            capture_test(self.keploy, req, res)
+
+        return resp

keploy/contrib/flask/utils.py

@@ -0,0 +1,20 @@
+
+def get_request_data(request) -> dict:
+    req_data = {}
+
+    req_data['header'] = {key: [value] for key,value in request.headers.to_wsgi_list()}
+    req_data['method'] = request.method
+    req_data['body'] = request.get_data(as_text=True)
+    # req_data['form_data'] = request.form.to_dict()
+    # req_data['file_data'] =  { k: v[0].read() for k, v in request.files.lists()}
+    req_data['uri'] = request.url_rule.rule
+    req_data['url'] = request.path
+    req_data['base'] = request.url
+    req_data['params'] = request.args.to_dict()
+
+    protocol = request.environ.get('SERVER_PROTOCOL', None)
+    if protocol:
+        req_data['proto_major'] = int(protocol.split(".")[0][-1])
+        req_data['proto_minor'] = int(protocol.split(".")[1])
+
+    return req_data

keploy/keploy.py

@@ -2,9 +2,10 @@
 import logging
 import http.client
 import re
+import threading
 import time
 from typing import Iterable, List, Mapping, Optional, Sequence
-from keploy.mode import mode
+from keploy.mode import getMode
 from keploy.constants import MODE_TEST, USE_HTTPS
 
 from keploy.models import Config, Dependency, HttpResp, TestCase, TestCaseRequest, TestReq
@@ -22,19 +23,20 @@ def __init__(self, conf:Config) -> None:
 
         self._config = conf
         self._logger = logger
-        self._dependencies:Mapping[str, List[Dependency]] = {}
-        self._responses:Mapping[str, HttpResp] = {}
+        self._dependencies = {}
+        self._responses = {}
         self._client = None
 
         if self._config.server.protocol == USE_HTTPS:
             self._client =  http.client.HTTPSConnection(host=self._config.server.url, port=self._config.server.port)
         else:
             self._client = http.client.HTTPConnection(host=self._config.server.url, port=self._config.server.port)
 
-        self._client.connect()
+        # self._client.connect()
 
-        if mode == MODE_TEST:
-            self.test()
+        if getMode() == MODE_TEST:
+            t = threading.Thread(target=self.test)
+            t.start()
 
 
     def get_dependencies(self, id: str) -> Optional[Iterable[Dependency]]:
@@ -96,7 +98,7 @@ def put(self, rq: TestCaseRequest):
                     return None
 
             headers = {'Content-type': 'application/json', 'key': self._config.server.licenseKey}
-            bytes_data = json.dumps(rq.__dict__).encode()
+            bytes_data = json.dumps(rq, default=lambda o: o.__dict__).encode()
             self._client.request("POST", "/{}/regression/testcase".format(self._config.server.suffix), bytes_data, headers)
 
             response = self._client.getresponse()
@@ -106,26 +108,29 @@ def put(self, rq: TestCaseRequest):
             body = json.loads(response.read().decode())
             if body.get('id', None):
                 self.denoise(body['id'], rq)
+
         except:
             self._logger.error("Exception occured while storing the request information. Try again.")
 
 
     def denoise(self, id:str, tcase:TestCaseRequest):
-        time.sleep(2.0)
         try:
-            unit = TestCase(id, captured=tcase.captured, uri=tcase.uri, req=tcase.httpRequest, deps=tcase.deps)
+            unit = TestCase(id=id, captured=tcase.captured, uri=tcase.uri, http_req=tcase.http_req, http_resp={}, deps=tcase.deps)
             res = self.simulate(unit)
             if not res:
-                self._logger.error("failed to simulate request")
+                self._logger.error("failed to simulate request while denoising")
                 return
 
             headers = {'Content-type': 'application/json', 'key': self._config.server.licenseKey}
-            bin_data = json.dumps(TestReq(id=id, app_id=self._config.app.name, resp=res).__dict__).encode()
+            bin_data = json.dumps(TestReq(id=id, app_id=self._config.app.name, resp=res), default=lambda o: o.__dict__).encode()
             self._client.request("POST", "/{}/regression/denoise".format(self._config.server.suffix), bin_data, headers)
 
             response = self._client.getresponse()
             if response.status != 200:
                 self._logger.error("failed to de-noise request to backend")
+
+            body = response.read().decode()
+
         except:
             self._logger.error("Error occured while denoising the test case request. Skipping...")
 
@@ -135,7 +140,7 @@ def simulate(self, test_case:TestCase) -> Optional[HttpResp]:
             self._dependencies[test_case.id] = test_case.deps
 
             heads = test_case.http_req.header
-            heads['KEPLOY_TEST_ID'] = test_case.id
+            heads['KEPLOY_TEST_ID'] = [test_case.id]
 
             cli = http.client.HTTPConnection(self._config.app.host, self._config.app.port)
             cli._http_vsn = int(str(test_case.http_req.proto_major) + str(test_case.http_req.proto_minor))
@@ -145,9 +150,11 @@ def simulate(self, test_case:TestCase) -> Optional[HttpResp]:
                 method=test_case.http_req.method,
                 url=self._config.app.suffix + test_case.http_req.url,
                 body=json.dumps(test_case.http_req.body).encode(),
-                headers=heads
+                headers={key: value[0] for key, value in heads.items()}
             )
 
+            time.sleep(2.0)
+            # TODO: getting None in case of regular execution. Urgent fix needed.
             response = self.get_resp(test_case.id)
             if not response or not self._responses.pop(test_case.id, None):
                 self._logger.error("failed loading the response for testcase.")
@@ -158,7 +165,7 @@ def simulate(self, test_case:TestCase) -> Optional[HttpResp]:
 
             return response
 
-        except  Exception as e:
+        except Exception as e:
             self._logger.exception("Exception occured in simulation of test case with id: %s" %test_case.id)
 
 
@@ -170,7 +177,7 @@ def check(self, r_id:str, tc: TestCase) -> bool:
                 return False
 
             headers = {'Content-type': 'application/json', 'key': self._config.server.licenseKey}
-            bytes_data = json.dumps(TestReq(id=tc.id, app_id=self._config.app.name, run_id=r_id, resp=resp).__dict__).encode()
+            bytes_data = json.dumps(TestReq(id=tc.id, app_id=self._config.app.name, run_id=r_id, resp=resp), default=lambda o: o.__dict__).encode()
             self._client.request("POST", "/{}/regression/test".format(self._config.server.suffix), bytes_data, headers)
 
             response = self._client.getresponse()
@@ -236,7 +243,6 @@ def fetch(self, offset:int=0, limit:int=25) -> Optional[Sequence[TestCase]]:
 
     def test(self):
         passed = True
-
         time.sleep(self._config.app.delay)
 
         self._logger.info("Started test operations on the captured test cases.")

keploy/mode.py

@@ -1,7 +1,8 @@
+from typing import Literal
 from keploy.constants import MODE_OFF, MODE_RECORD, MODE_TEST
 
 
-mode = MODE_OFF
+mode = MODE_RECORD
 
 def isValidMode(mode):
 	if mode in [MODE_OFF, MODE_RECORD, MODE_TEST]:
@@ -11,7 +12,8 @@ def isValidMode(mode):
 def getMode():
 	return mode
 
-def setMode(m):
+MODES = Literal["off", "record", "test"]
+def setMode(m:MODES):
 	if isValidMode(m):
 		global mode
 		mode = m

keploy/models.py

@@ -113,7 +113,7 @@ def __init__(self, method:METHODS=None, proto_major:int=1, proto_minor:int=1, ur
 
 class HttpResp(object):
 	def __init__(self, status_code:int=None, header:Mapping[str, Sequence[str]]=None, body:str=None) -> None:
-		self.code = status_code
+		self.status_code = status_code
 		self.header = header
 		self.body = body
 
@@ -123,28 +123,39 @@ def __init__(self, id:str=None, created:int=None, updated:int=None, captured:int
 				 cid:str=None, app_id:str=None, uri:str=None, http_req:HttpReq=None,
 				 http_resp: HttpResp=None, deps:Sequence[Dependency]=None, all_keys:Mapping[str, Sequence[str]]=None,
 				 anchors:Mapping[str, Sequence[str]]=None, noise:Sequence[str]=None ) -> None:
+
+		#TODO: Need to handle the case when the API response is None instead of [] for deps
+		if not deps:
+			deps = []
+
 		self.id = id
 		self.created = created
 		self.updated = updated
 		self.captured = captured
-		self.c_id = cid
+		self.cid = cid
 		self.app_id = app_id
 		self.uri = uri
-		self.http_req = HttpReq(**http_req)
-		self.http_resp = HttpResp(**http_resp)
-		self.deps = [Dependency(**dep) for dep in deps]
+		self.http_req = http_req
+		self.http_resp = http_resp
+		self.deps = [Dependency(**dep) if not isinstance(dep, Dependency) else dep for dep in deps]
 		self.all_keys = all_keys
 		self.anchors = anchors
 		self.noise = noise
 
+		if not isinstance(http_req, HttpReq):
+			self.http_req = HttpReq(**http_req)
+
+		if not isinstance(http_resp, HttpResp):
+			self.http_resp = HttpResp(**http_resp)
+
 
 class TestCaseRequest(object):
 	def __init__(self, captured:int=None, app_id:str=None, uri:str=None, http_req:HttpReq=None, http_resp:HttpResp=None, deps:Iterable[Dependency]=None) -> None:
 		self.captured = captured
 		self.app_id = app_id
 		self.uri = uri
-		self.httpRequest = http_req
-		self.httpResponse = http_resp
+		self.http_req = http_req
+		self.http_resp = http_resp
 		self.deps = deps
 
 		if not captured:

keploy/utils.py

@@ -0,0 +1,32 @@
+from keploy.keploy import Keploy
+from keploy.models import Dependency, HttpReq, HttpResp, TestCaseRequest
+import time
+
+def capture_test(k, reqst, resp):
+
+    deps = [ Dependency('demo_dep', 'HTTP_CLIENT', {}, None), ]
+
+    test = TestCaseRequest(
+        captured=int(time.time()),
+        app_id=k._config.app.name,
+        uri=reqst['uri'],
+        http_req=HttpReq(
+            method=reqst['method'],
+            proto_major=reqst['proto_major'],
+            proto_minor=reqst['proto_minor'],
+            url=reqst['url'],
+            url_params=reqst['params'],
+            header=reqst['header'],
+            body=reqst['body']
+        ),
+        http_resp=HttpResp(
+            status_code=resp['status_code'],
+            header=resp['header'],
+            body=resp['body']
+        ),
+        deps=deps
+    )
+
+    k.capture(test)
+
+    return