Fix/analytics connection hang

src/mlstacks/analytics/client.py

@@ -16,7 +16,7 @@
 import os
 from logging import getLogger
 from types import TracebackType
-from typing import Any, Dict, Optional, Type, cast
+from typing import Any, Dict, List, Optional, Type, cast
 from uuid import uuid4
 
 import click
@@ -34,10 +34,24 @@
 logger = getLogger(__name__)
 
 analytics.write_key = "tU9BJvF05TgC29xgiXuKF7CuYP0zhgnx"
+analytics.max_retries = 1
 
 CONFIG_FILENAME = "config.yaml"
 
 
+def on_error(error: Exception, batch: List[Dict[str, Any]]) -> None:
+    """Custom error handler for Segment analytics.
+
+    Args:
+        error: The error that occurred.
+        batch: Events processed when the error occurred.
+    """
+    logger.debug("Analytics error: %s; Batch: %s", error, batch)
+
+
+analytics.on_error = on_error
+
+
 class MLStacksAnalyticsContext:
     """Analytics context manager for MLStacks."""
 
@@ -161,7 +175,7 @@ def track_event(
         metadata: Dict of metadata to track.
 
     Returns:
-        True if event is sent successfully, False is not.
+        True if event is sent successfully, False otherwise.
     """
     if metadata is None:
         metadata = {}

tests/unit/analytics/__init__.py

@@ -0,0 +1,12 @@
+#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.

tests/unit/analytics/test_custom_on_error.py

@@ -0,0 +1,59 @@
+#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+import logging
+from typing import Any, Dict, List
+from unittest.mock import patch
+
+import pytest
+import requests
+from segment import analytics
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.fixture(autouse=True)
+def reset_analytics_client():
+    default_on_error = analytics.on_error
+
+    yield
+
+    analytics.on_error = default_on_error
+
+
+def mock_post_failure(*args, **kwargs):
+    response = requests.Response()
+    response.status_code = 500
+    response._content = (
+        b'{"error": "Simulated failure", "code": "internal_error"}'
+    )
+    response.headers["Content-Type"] = "application/json"
+    return response
+
+
+def custom_on_error(error: Exception, batch: List[Dict[str, Any]]) -> None:
+    logger.debug("Analytics error: %s; Batch: %s", error, batch)
+
+
+@pytest.mark.usefixtures("reset_analytics_client")
+def test_segment_custom_on_error_handler_invocation(caplog):
+    with caplog.at_level(logging.DEBUG):
+        with patch(
+            "segment.analytics.request._session.post",
+            side_effect=mock_post_failure,
+        ):
+            analytics.track(
+                "test_user_id", "Test Event", {"property": "value"}
+            )
+            analytics.flush()
+
+    assert "Analytics error:" in caplog.text