Merge pull request #21 from ISISComputingGroup/error_handling_and_tests

Error handling and tests

Author: Tom-Willemsen

.gitignore

@@ -50,6 +50,7 @@ coverage.xml
 .hypothesis/
 .pytest_cache/
 cover/
+coverage_html_report/
 
 # Translations
 *.mo

pyproject.toml

@@ -7,7 +7,8 @@ name = "saluki"
 dynamic = ["version"]
 dependencies = [
     "ess-streaming-data-types",
-    "confluent-kafka"
+    "confluent-kafka",
+    "tzdata"
 ]
 readme = {file = "README.md", content-type = "text/markdown"}
 license-files = ["LICENSE"]
@@ -25,6 +26,29 @@ dev = [
     "pyright",
     "ruff",
     "pytest",
+    "pytest-cov"
 ]
 
+[tool.coverage.run]
+branch = true
+source = ["src"]
+
+[tool.coverage.report]
+fail_under = 100
+exclude_lines = [
+  "pragma: no cover",
+  "if TYPE_CHECKING:",
+  "if typing.TYPE_CHECKING:",
+  "@abstractmethod",
+]
+omit = ["main.py"]  # No logic here besides argparse
+
+
+[tool.coverage.html]
+directory = "coverage_html_report"
+
+[tool.pytest.ini_options]
+testpaths = "tests"
+addopts = "--cov --cov-report=html -vv"
+
 [tool.setuptools_scm]

src/saluki/__init__.py

@@ -1,60 +1 @@
-import logging
-from typing import Tuple, List
-import datetime
-from confluent_kafka import Message
-from streaming_data_types import DESERIALISERS
-from streaming_data_types.exceptions import StreamingDataTypesException
-from streaming_data_types.utils import get_schema
 
-logger = logging.getLogger("saluki")
-
-
-def _fallback_deserialiser(payload: bytes) -> str:
-    return payload.decode()
-
-
-def try_to_deserialise_message(payload: bytes) -> Tuple[str, str]:
-    logger.debug(f"got some data: {payload}")
-    schema = get_schema(payload)
-    deserialiser = (
-        _fallback_deserialiser  # Fall back to this if we need to so data isn't lost
-    )
-    try:
-        deserialiser = DESERIALISERS[schema]
-    except StreamingDataTypesException:
-        pass  # TODO
-    except KeyError:
-        pass
-    return schema, deserialiser(payload)
-
-
-def _deserialise_and_print_messages(msgs: List[Message], partition: int | None) -> None:
-    for msg in msgs:
-        if msg is None:
-            continue
-        if msg.error():
-            logger.error("Consumer error: {}".format(msg.error()))
-            continue
-        if partition is not None and msg.partition() != partition:
-            continue
-        schema, deserialised = try_to_deserialise_message(msg.value())
-        time = _parse_timestamp(msg)
-        logger.info(f"{msg.offset()} ({time}):({schema}) {deserialised}")
-
-
-def _parse_timestamp(msg: Message) -> str:
-    """
-    Parse a message timestamp.
-
-    See https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#confluent_kafka.Message.timestamp
-    :param msg: the message to parse.
-    :return: either the string-formatted timestamp or "Unknown" if not able to parse.
-    """
-    timestamp_type, timestamp_ms_from_epoch = msg.timestamp()
-    if timestamp_type == 1:  # TIMESTAMP_CREATE_TIME
-        return datetime.datetime.fromtimestamp(timestamp_ms_from_epoch / 1000).strftime(
-            "%Y-%m-%d %H:%M:%S.%f"
-        )
-    else:
-        # TIMESTAMP_NOT_AVAILABLE or TIMESTAMP_LOG_APPEND_TIME
-        return "Unknown"

src/saluki/consume.py

@@ -1,7 +1,8 @@
 import logging
 
 from confluent_kafka import Consumer, TopicPartition
-from saluki import _deserialise_and_print_messages
+
+from saluki.utils import deserialise_and_print_messages
 
 logger = logging.getLogger("saluki")
 
@@ -39,29 +40,26 @@ def consume(
 
     if go_forwards:
         if offset is None:
-            logger.error("Can't go forwards without an offset")
-            return
+            raise ValueError("Can't go forwards without an offset")
         start = offset
     else:
         if offset is not None:
             start = offset - num_messages + 1
         else:
             start = (
-                c.get_watermark_offsets(TopicPartition(topic, partition), cached=False)[
-                    1
-                ]
+                c.get_watermark_offsets(TopicPartition(topic, partition), cached=False)[1]
                 - num_messages
             )
 
-    logger.info(f"starting at offset {start}")
+    logger.info(f"Starting at offset {start}")
     c.assign([TopicPartition(topic, partition, start)])
 
     try:
-        logger.info(f"consuming {num_messages} messages")
+        logger.info(f"Consuming {num_messages} messages")
         msgs = c.consume(num_messages)
-        _deserialise_and_print_messages(msgs, partition)
-    except Exception as e:
-        logger.error(e)
+        deserialise_and_print_messages(msgs, partition)
+    except Exception:
+        logger.exception("Got exception while consuming:")
     finally:
-        logger.debug(f"closing consumer {c}")
+        logger.debug(f"Closing consumer {c}")
         c.close()

src/saluki/listen.py

@@ -1,7 +1,8 @@
 import logging
 
 from confluent_kafka import Consumer, TopicPartition
-from saluki import _deserialise_and_print_messages
+
+from saluki.utils import deserialise_and_print_messages
 
 logger = logging.getLogger("saluki")
 
@@ -29,7 +30,7 @@ def listen(broker: str, topic: str, partition: int | None = None) -> None:
         logger.info(f"listening to {broker}/{topic}")
         while True:
             msg = c.poll(1.0)
-            _deserialise_and_print_messages([msg], partition)
+            deserialise_and_print_messages([msg], partition)
     except KeyboardInterrupt:
         logger.debug("finished listening")
     finally:

src/saluki/main.py

@@ -1,10 +1,10 @@
 import argparse
 import logging
 import sys
-from typing import Tuple
 
 from saluki.consume import consume
 from saluki.listen import listen
+from saluki.utils import parse_kafka_uri
 
 logger = logging.getLogger("saluki")
 logging.basicConfig(level=logging.INFO)
@@ -14,30 +14,6 @@
 _CONSUME = "consume"
 
 
-def parse_kafka_uri(uri: str) -> Tuple[str, str]:
-    """Parse Kafka connection URI.
-
-    A broker hostname/ip must be present.
-    If username is provided, a SASL mechanism must also be provided.
-    Any other validation must be performed in the calling code.
-    """
-    security_protocol, tail = uri.split("+") if "+" in uri else ("", uri)
-    sasl_mechanism, tail = tail.split("\\") if "\\" in tail else ("", tail)
-    username, tail = tail.split("@") if "@" in tail else ("", tail)
-    broker, topic = tail.split("/") if "/" in tail else (tail, "")
-    if not broker:
-        raise RuntimeError(
-            f"Unable to parse URI {uri}, broker not defined. URI should be of form"
-            f" [PROTOCOL+SASL_MECHANISM\\username@]broker:9092"
-        )
-    if username and not (security_protocol and sasl_mechanism):
-        raise RuntimeError(
-            f"Unable to parse URI {uri}, PROTOCOL or SASL_MECHANISM not defined."
-            f" URI should be of form [PROTOCOL+SASL_MECHANISM\\username@]broker:9092"
-        )
-    return broker, topic
-
-
 def main() -> None:
     parser = argparse.ArgumentParser(
         prog="saluki",

src/saluki/utils.py

@@ -0,0 +1,87 @@
+import datetime
+import logging
+from typing import List, Tuple
+from zoneinfo import ZoneInfo
+
+from confluent_kafka import Message
+from streaming_data_types import DESERIALISERS
+from streaming_data_types.exceptions import ShortBufferException
+from streaming_data_types.utils import get_schema
+
+logger = logging.getLogger("saluki")
+
+
+def _try_to_deserialise_message(payload: bytes) -> Tuple[str | None, str | None]:
+    logger.debug(f"got some data: {payload}")
+    try:
+        schema = get_schema(payload)
+    except ShortBufferException:
+        schema = None
+
+    logger.debug(f"schema: {schema}")
+
+    def fallback_deserialiser(payload: bytes) -> str:
+        return payload.decode()
+
+    deserialiser = DESERIALISERS.get(schema if schema is not None else "", fallback_deserialiser)
+    logger.debug(f"Deserialiser: {deserialiser}")
+
+    ret = deserialiser(payload)
+
+    return schema, ret
+
+
+def deserialise_and_print_messages(msgs: List[Message], partition: int | None) -> None:
+    for msg in msgs:
+        try:
+            if msg is None:
+                continue
+            if msg.error():
+                logger.error("Consumer error: {}".format(msg.error()))
+                continue
+            if partition is not None and msg.partition() != partition:
+                continue
+            schema, deserialised = _try_to_deserialise_message(msg.value())
+            time = _parse_timestamp(msg)
+            logger.info(f"{msg.offset()} ({time}):({schema}) {deserialised}")
+        except Exception as e:
+            logger.exception(f"Got error while deserialising: {e}")
+
+
+def _parse_timestamp(msg: Message) -> str:
+    """
+    Parse a message timestamp.
+
+    See https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#confluent_kafka.Message.timestamp
+    :param msg: the message to parse.
+    :return: either the string-formatted timestamp or "Unknown" if not able to parse.
+    """
+    timestamp_type, timestamp_ms_from_epoch = msg.timestamp()
+    if timestamp_type == 1:  # TIMESTAMP_CREATE_TIME
+        return (
+            datetime.datetime.fromtimestamp(timestamp_ms_from_epoch / 1000)
+            .astimezone(ZoneInfo("UTC"))
+            .strftime("%Y-%m-%d %H:%M:%S.%f")
+        )
+    else:
+        # TIMESTAMP_NOT_AVAILABLE or TIMESTAMP_LOG_APPEND_TIME
+        return "Unknown"
+
+
+def parse_kafka_uri(uri: str) -> Tuple[str, str]:
+    """Parse Kafka connection URI.
+
+    A broker hostname/ip must be present.
+    If username is provided, a SASL mechanism must also be provided.
+    Any other validation must be performed in the calling code.
+    """
+    broker, topic = uri.split("/") if "/" in uri else (uri, "")
+    if not topic:
+        raise RuntimeError(
+            f"Unable to parse URI {uri}, topic not defined. URI should be of form"
+            f" broker[:port]/topic"
+        )
+    return (
+        broker,
+        topic,
+    )

tests/test_consume.py

@@ -0,0 +1,84 @@
+from unittest import mock
+from unittest.mock import patch
+
+import pytest
+from confluent_kafka import TopicPartition
+
+from saluki.consume import consume
+
+
+@patch("saluki.consume.Consumer")
+def test_go_forwards_with_no_offset_raises(_):
+    with pytest.raises(ValueError):
+        consume("broker", "topic", go_forwards=True, offset=None)
+
+
+@patch("saluki.consume.Consumer")
+def test_go_forwards_with_offset_assigns_at_offset(mock_consumer):
+    expected_topic = "topic"
+    expected_offset = 1234
+    expected_partition = 1
+    consume(
+        "broker",
+        expected_topic,
+        go_forwards=True,
+        offset=expected_offset,
+        partition=expected_partition,
+    )
+    mock_assign = mock_consumer.return_value.assign
+
+    mock_assign.assert_called_with(
+        [TopicPartition(expected_topic, expected_partition, expected_offset)]
+    )
+
+
+@patch("saluki.consume.Consumer")
+def test_consume_with_offset_and_num_of_messages_goes_back_offset_minus_messages(
+    mock_consumer,
+):
+    expected_offset = 1234
+    expected_topic = "sometopic"
+    num_messages = 3
+    expected_start_offset = expected_offset - num_messages + 1
+
+    consume("broker", expected_topic, offset=expected_offset, num_messages=num_messages)
+
+    mock_assign = mock_consumer.return_value.assign
+    mock_assign.assert_called_once()
+
+    mock_assign_call = mock_assign.call_args.args[0][0]
+    assert mock_assign_call.topic == expected_topic
+    assert mock_assign_call.offset == expected_start_offset
+
+
+@patch("saluki.consume.Consumer")
+def test_consume_with_no_offset_and_num_of_messages_goes_back_high_watermark_minus_messages(
+    mock_consumer,
+):
+    expected_topic = "sometopic"
+    num_messages = 3
+    high_watermark_offset = 2345
+    expected_start_offset = high_watermark_offset - num_messages
+
+    mock_consumer.return_value.get_watermark_offsets.return_value = (
+        None,
+        high_watermark_offset,
+    )
+
+    consume("broker", topic=expected_topic, num_messages=num_messages)
+    mock_assign = mock_consumer.return_value.assign
+    mock_assign.assert_called_once()
+
+    mock_assign_call = mock_assign.call_args.args[0][0]
+
+    assert mock_assign_call.topic == expected_topic
+    assert mock_assign_call.offset == expected_start_offset
+
+
+def test_consume_but_exception_thrown_consumer_is_closed():
+    with (
+        mock.patch("saluki.consume.Consumer") as c,
+    ):
+        c.return_value.consume.side_effect = Exception
+        consume("somebroker", "sometopic", num_messages=1)
+        c.return_value.close.assert_called_once()