Module: check_mk
Branch: master
Commit: 5340cec1b2e0ae5ef68e066eb205fccea1121933
URL:
http://git.mathias-kettner.de/git/?p=check_mk.git;a=commit;h=5340cec1b2e0ae…
Author: Óscar Nájera <on(a)mathias-kettner.de>
Date: Thu Feb 21 15:59:09 2019 +0100
A bit more py3 compatible livestatus
Change-Id: I0a6df3acf69e860adb6d6d98909f12c9b817f0c9
---
livestatus/api/python/livestatus.py | 31 ++++++++++++++-------------
tests/unit/livestatus/test_livestatus_unit.py | 10 ++++++---
2 files changed, 23 insertions(+), 18 deletions(-)
diff --git a/livestatus/api/python/livestatus.py b/livestatus/api/python/livestatus.py
index 20ee562..d36e6d3 100644
--- a/livestatus/api/python/livestatus.py
+++ b/livestatus/api/python/livestatus.py
@@ -52,6 +52,12 @@ persistent_connections = {} # type: Dict[str, socket.socket]
remove_cache_regex = re.compile("\nCache:[^\n]*") # type: Pattern
+def ensure_unicode(text):
+ if hasattr(text, "decode"):
+ return text.decode("utf-8")
+ return text
+
+
class MKLivestatusException(Exception):
def __init__(self, value):
self.parameter = value
@@ -97,7 +103,7 @@ def lqencode(s):
# It is not enough to strip off \n\n, because one might submit "\n \n",
# which is also interpreted as termination of the last query and beginning
# of the next query.
- return s.replace('\n', '')
+ return ensure_unicode(s).replace(u"\n", u"")
def site_local_ca_path():
@@ -238,19 +244,13 @@ class Query(object):
def __init__(self, query, suppress_exceptions=None):
super(Query, self).__init__()
- self._query = self._ensure_unicode(query)
+ self._query = ensure_unicode(query)
if suppress_exceptions is None:
self.suppress_exceptions = self.default_suppressed_exceptions
else:
self.suppress_exceptions = suppress_exceptions
- def _ensure_unicode(self, thing):
- try:
- return unicode(thing)
- except UnicodeDecodeError:
- return thing.decode("utf-8")
-
def __unicode__(self):
return self._query
@@ -377,7 +377,7 @@ class SingleSiteConnection(Helpers):
if family_txt == "unix":
return socket.AF_UNIX, url
- elif family_txt in ["tcp", "tcp6"]:
+ if family_txt in ["tcp", "tcp6"]:
try:
host, port_txt = url.rsplit(":", 1)
port = int(port_txt)
@@ -407,16 +407,16 @@ class SingleSiteConnection(Helpers):
pass
def receive_data(self, size):
- result = ""
+ result = u""
# Timeout is only honored when connecting
self.socket.settimeout(None)
while size > 0:
packet = self.socket.recv(size)
- if len(packet) == 0:
+ if not packet:
raise MKLivestatusSocketClosed(
"Read zero data from socket, nagios server closed
connection")
size -= len(packet)
- result += packet
+ result += packet.decode("utf-8")
return result
def do_query(self, query, add_headers=""):
@@ -447,8 +447,7 @@ class SingleSiteConnection(Helpers):
try:
# socket.send() will implicitely cast to str(), we need ot
# convert to UTF-8 in order to avoid exceptions
- if isinstance(query, unicode):
- query = query.encode("utf-8")
+ query = query.encode("utf-8")
self.socket.send(query)
except IOError as e:
if self.persist:
@@ -641,7 +640,9 @@ class MultiSiteConnection(Helpers):
# hosts at the same time.
status_hosts = {} # dict from site to list of status_hosts
- for sitename, site in sites.items() + extra_status_sites.items():
+ sites_dict = sites.copy()
+ sites_dict.update(extra_status_sites)
+ for sitename, site in sites_dict.items():
status_host = site.get("status_host")
if status_host:
if not isinstance(status_host, tuple) or len(status_host) != 2:
diff --git a/tests/unit/livestatus/test_livestatus_unit.py
b/tests/unit/livestatus/test_livestatus_unit.py
index c28002f..24435e4 100644
--- a/tests/unit/livestatus/test_livestatus_unit.py
+++ b/tests/unit/livestatus/test_livestatus_unit.py
@@ -1,7 +1,11 @@
import socket
import ssl
from contextlib import closing
-from pathlib2 import Path
+try:
+ from pathlib import Path # Py3 first
+except ImportError:
+ from pathlib2 import Path
+
import pytest # type: ignore
import omdlib.certs as certs
@@ -25,13 +29,13 @@ def sock_path(monkeypatch, tmpdir):
@pytest.mark.parametrize("query_part", [
+ "xyz\nabc",
u"xyz\nabc",
b"xyz\nabc",
])
def test_lqencode(query_part):
result = livestatus.lqencode(query_part)
- assert isinstance(result, type(query_part))
- assert result == "xyzabc"
+ assert result == u"xyzabc"
def test_livestatus_local_connection_omd_root_not_set(monkeypatch, tmpdir):