Injectable statsd client

requirements-dev.txt

@@ -15,6 +15,7 @@
 # limitations under the License.
 #
 console_log==0.2.10
+coverage==4.5.3
 flake8-commas==2.0.0
 flake8-import-order==0.18
 flake8-mypy==17.8.0
@@ -24,6 +25,7 @@ flask-cors==3.0.6
 ipdb==0.11
 mypy==0.670
 mysqlclient==1.3.13
+nose==1.3.7
 pip-tools==3.5.0
 psycopg2-binary==2.7.5
 pycodestyle==2.4.0

superset/stats_logger.py

@@ -73,8 +73,19 @@ def gauge(self, key, value):
     from statsd import StatsClient
 
     class StatsdStatsLogger(BaseStatsLogger):
-        def __init__(self, host, port, prefix='superset'):
-            self.client = StatsClient(host=host, port=port, prefix=prefix)
+
+        def __init__(self, host='localhost', port=8125,
+                     prefix='superset', statsd_client=None):
+            """
+            Initializes from either params or a supplied, pre-constructed statsd client.
+
+            If statsd_client argument is given, all other arguments are ignored and the
+            supplied client will be used to emit metrics.
+            """
+            if statsd_client:
+                self.client = statsd_client
+            else:
+                self.client = StatsClient(host=host, port=port, prefix=prefix)
 
         def incr(self, key):
             self.client.incr(key)

tests/access_tests.py

@@ -17,8 +17,7 @@
 """Unit tests for Superset"""
 import json
 import unittest
-
-import mock
+from unittest import mock
 
 from superset import app, db, security_manager
 from superset.connectors.connector_registry import ConnectorRegistry

tests/base_tests.py

@@ -17,9 +17,9 @@
 """Unit tests for Superset"""
 import json
 import unittest
+from unittest.mock import Mock, patch
 
 from flask_appbuilder.security.sqla import models as ab_models
-from mock import Mock, patch
 import pandas as pd
 
 from superset import app, db, is_feature_enabled, security_manager

tests/core_tests.py

@@ -26,8 +26,8 @@
 import re
 import string
 import unittest
+from unittest import mock
 
-import mock
 import pandas as pd
 import psycopg2
 import sqlalchemy as sqla

tests/db_engine_specs_test.py

@@ -15,8 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 import inspect
+from unittest import mock
 
-import mock
 from sqlalchemy import column, select, table
 from sqlalchemy.dialects.mssql import pymssql
 from sqlalchemy.types import String, UnicodeText

tests/druid_func_tests.py

@@ -16,8 +16,8 @@
 # under the License.
 import json
 import unittest
+from unittest.mock import Mock
 
-from mock import Mock
 from pydruid.utils.dimensions import MapLookupExtraction, RegexExtraction
 import pydruid.utils.postaggregator as postaggs
 

tests/druid_tests.py

@@ -18,8 +18,7 @@
 from datetime import datetime
 import json
 import unittest
-
-from mock import Mock, patch
+from unittest.mock import Mock, patch
 
 from superset import db, security_manager
 from superset.connectors.druid.models import (

tests/email_tests.py

@@ -22,8 +22,7 @@
 import logging
 import tempfile
 import unittest
-
-import mock
+from unittest import mock
 
 from superset import app
 from superset.utils import core as utils

tests/schedules_test.py

@@ -16,9 +16,9 @@
 # under the License.
 from datetime import datetime, timedelta
 import unittest
+from unittest.mock import Mock, patch, PropertyMock
 
 from flask_babel import gettext as __
-from mock import Mock, patch, PropertyMock
 from selenium.common.exceptions import WebDriverException
 
 from superset import app, db

tests/stats_logger_tests.py

@@ -0,0 +1,34 @@
+from unittest import TestCase
+from unittest.mock import Mock, patch
+
+from superset.stats_logger import StatsdStatsLogger
+
+
+class StatsdStatsLoggerTest(TestCase):
+
+    def verify_client_calls(self, logger, client):
+        logger.incr('foo1')
+        client.incr.assert_called_once()
+        client.incr.assert_called_with('foo1')
+        logger.decr('foo2')
+        client.decr.assert_called_once()
+        client.decr.assert_called_with('foo2')
+        logger.gauge('foo3')
+        client.gauge.assert_called_once()
+        client.gauge.assert_called_with('foo3')
+        logger.timing('foo4', 1.234)
+        client.timing.assert_called_once()
+        client.timing.assert_called_with('foo4', 1.234)
+
+    def test_init_with_statsd_client(self):
+        client = Mock()
+        stats_logger = StatsdStatsLogger(statsd_client=client)
+        self.verify_client_calls(stats_logger, client)
+
+    def test_init_with_params(self):
+        with patch('superset.stats_logger.StatsClient') as MockStatsdClient:
+            mock_client = MockStatsdClient.return_value
+
+            stats_logger = StatsdStatsLogger()
+            self.verify_client_calls(stats_logger, mock_client)
+

tests/utils_tests.py

@@ -17,9 +17,9 @@
 from datetime import date, datetime, time, timedelta
 from decimal import Decimal
 import unittest
+from unittest.mock import patch
 import uuid
 
-from mock import patch
 import numpy
 
 from superset.exceptions import SupersetException

tests/viz_tests.py

@@ -15,9 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 from datetime import datetime
+from unittest.mock import Mock, patch
 import uuid
 
-from mock import Mock, patch
 import pandas as pd
 
 from superset import app

tox.ini

@@ -51,9 +51,6 @@ commands =
 deps =
     -rrequirements.txt
     -rrequirements-dev.txt
-    coverage
-    mock
-    nose
 setenv =
     PYTHONPATH = {toxinidir}
     SUPERSET_CONFIG = tests.superset_test_config