From ae42f2d55917fe0b9f02b37c83872f59b1505bd2 Mon Sep 17 00:00:00 2001
From: Anton Sidelnikov <53078276+anton-sidelnikov@users.noreply.github.com>
Date: Tue, 2 Feb 2021 15:42:19 +0300
Subject: [PATCH] LB load new util (#42)

LB load new util

Reviewed-by: Anton Kachurin <katchuring@gmail.com>
Reviewed-by: Rodion Gyrbu <fpsoff@outlook.com>
Reviewed-by: Anton Sidelnikov <None>
Reviewed-by: None <None>
---
 csm_test_utils/__main__.py                |  3 ++
 csm_test_utils/loadbalancer/__init__.py   |  1 +
 csm_test_utils/loadbalancer/lb_monitor.py | 44 +++++++++++++++++++++++
 3 files changed, 48 insertions(+)
 create mode 100644 csm_test_utils/loadbalancer/__init__.py
 create mode 100644 csm_test_utils/loadbalancer/lb_monitor.py

diff --git a/csm_test_utils/__main__.py b/csm_test_utils/__main__.py
index 3aea7c2..c764930 100644
--- a/csm_test_utils/__main__.py
+++ b/csm_test_utils/__main__.py
@@ -6,6 +6,7 @@
 from csm_test_utils.rebalance_test import main as r_main
 from csm_test_utils.dns import dns_lookup_main, host_main
 from csm_test_utils.rds import rds_backup_main
+from csm_test_utils.loadbalancer import lb_monitor
 
 args = root_parser.parse_args()
 if args.test == "monitor":
@@ -26,3 +27,5 @@
     host_main()
 if args.test == "rds_backup_monitor":
     rds_backup_main()
+if args.test == "lb_load":
+    lb_monitor()
diff --git a/csm_test_utils/loadbalancer/__init__.py b/csm_test_utils/loadbalancer/__init__.py
new file mode 100644
index 0000000..712eb2f
--- /dev/null
+++ b/csm_test_utils/loadbalancer/__init__.py
@@ -0,0 +1 @@
+from .lb_monitor import main as lb_monitor
diff --git a/csm_test_utils/loadbalancer/lb_monitor.py b/csm_test_utils/loadbalancer/lb_monitor.py
new file mode 100644
index 0000000..219c133
--- /dev/null
+++ b/csm_test_utils/loadbalancer/lb_monitor.py
@@ -0,0 +1,44 @@
+import logging
+import socket
+
+import requests
+from ocomone.logging import setup_logger
+
+from ..common import base_parser, sub_parsers
+
+LB_TIMING = "lb_timing"
+LB_TIMEOUT = "lb_timeout"
+
+LOGGER = logging.getLogger(__name__)
+LOGGER.setLevel(logging.DEBUG)
+
+AGP = sub_parsers.add_parser("lb_load", add_help=False, parents=[base_parser])
+
+
+def main():
+    args, _ = AGP.parse_known_args()
+    setup_logger(LOGGER, "lb_load", log_dir=args.log_dir, log_format="[%(asctime)s] %(message)s")
+    timeout = 20
+    try:
+        res = requests.get(args.target, headers={"Connection": "close"}, timeout=timeout)
+    except requests.Timeout as ex:
+        LOGGER.exception("Timeout sending request to LB")
+        result = {
+            "reason": LB_TIMEOUT,
+            "client": socket.gethostname(),
+            "timeout": timeout * 1000,
+            "exception": ex
+        }
+    else:
+        result = {
+            "reason": LB_TIMING,
+            "client": socket.gethostname(),
+            "server": res.headers["Server"],
+            "elapsed": res.elapsed.microseconds / 1000
+        }
+
+    print(result)
+
+
+if __name__ == "__main__":
+    main()