62 lines
2.2 KiB
Text
62 lines
2.2 KiB
Text
|
diff --git a/websocket_exporter/probe.py b/websocket_exporter/probe.py
|
||
|
|
index a95b97e..a7b057e 100644
|
||
|
|
--- a/websocket_exporter/probe.py
|
||
|
|
+++ b/websocket_exporter/probe.py
|
||
|
|
@@ -3,28 +3,29 @@ import logging
|
||
|
|
from time import perf_counter
|
||
|
|
from typing import Union
|
||
|
|
|
||
|
|
-from websockets import NegotiationError, client, InvalidStatusCode
|
||
|
|
+from websockets import InvalidStatusCode, NegotiationError, client
|
||
|
|
|
||
|
|
-
|
||
|
|
-EXACT_MATCH = 'exact'
|
||
|
|
-CONTAINS_MATCH = 'contains'
|
||
|
|
+EXACT_MATCH = "exact"
|
||
|
|
+CONTAINS_MATCH = "contains"
|
||
|
|
|
||
|
|
|
||
|
|
class ProbResults(object):
|
||
|
|
def __init__(self, up: int, latency: float = 0, received: int = 0):
|
||
|
|
self.up = up
|
||
|
|
self.latency = round(latency, 2)
|
||
|
|
- self.received = int(received) if received is not None else "NaN"
|
||
|
|
+ self.received = int(received) if received is not None else 0
|
||
|
|
|
||
|
|
def __str__(self):
|
||
|
|
if self.up:
|
||
|
|
return f'Websocket up, latency:{self.latency}s, expected response {"" if self.received else "NOT"} received'
|
||
|
|
- return f'Webserver DOWN'
|
||
|
|
+ return f"Webserver DOWN"
|
||
|
|
|
||
|
|
|
||
|
|
class WebSocketProbe(object):
|
||
|
|
|
||
|
|
- def __init__(self, uri, message=None, expected=None, match=CONTAINS_MATCH, timeout=10):
|
||
|
|
+ def __init__(
|
||
|
|
+ self, uri, message=None, expected=None, match=CONTAINS_MATCH, timeout=10
|
||
|
|
+ ):
|
||
|
|
"""
|
||
|
|
Create a websocket probe that tries establishing a connection and reports the metrics
|
||
|
|
:param uri: starts with 'ws://' or ws://
|
||
|
|
@@ -68,13 +69,17 @@ class WebSocketProbe(object):
|
||
|
|
elapsed = 0
|
||
|
|
while elapsed < self.timeout:
|
||
|
|
try:
|
||
|
|
- resp = await asyncio.wait_for(connection.recv(), timeout=(self.timeout-elapsed))
|
||
|
|
+ resp = await asyncio.wait_for(
|
||
|
|
+ connection.recv(), timeout=(self.timeout - elapsed)
|
||
|
|
+ )
|
||
|
|
if self._match(resp):
|
||
|
|
return True
|
||
|
|
await asyncio.sleep(1)
|
||
|
|
elapsed += 1
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
- logging.info(f'Time out while waiting for {self.expected_message} from {self.uri}')
|
||
|
|
+ logging.info(
|
||
|
|
+ f"Time out while waiting for {self.expected_message} from {self.uri}"
|
||
|
|
+ )
|
||
|
|
return None
|
||
|
|
return None
|
||
|
|
|