diff --git a/homeassistant/components/graphite/__init__.py b/homeassistant/components/graphite/__init__.py index ccb7044ad73..9405b576b4d 100644 --- a/homeassistant/components/graphite/__init__.py +++ b/homeassistant/components/graphite/__init__.py @@ -12,6 +12,7 @@ from homeassistant.const import ( CONF_HOST, CONF_PORT, CONF_PREFIX, + CONF_PROTOCOL, EVENT_HOMEASSISTANT_START, EVENT_HOMEASSISTANT_STOP, EVENT_STATE_CHANGED, @@ -21,8 +22,11 @@ import homeassistant.helpers.config_validation as cv _LOGGER = logging.getLogger(__name__) +PROTOCOL_TCP = "tcp" +PROTOCOL_UDP = "udp" DEFAULT_HOST = "localhost" DEFAULT_PORT = 2003 +DEFAULT_PROTOCOL = PROTOCOL_TCP DEFAULT_PREFIX = "ha" DOMAIN = "graphite" @@ -32,6 +36,9 @@ CONFIG_SCHEMA = vol.Schema( { vol.Optional(CONF_HOST, default=DEFAULT_HOST): cv.string, vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port, + vol.Optional(CONF_PROTOCOL, default=DEFAULT_PROTOCOL): vol.Any( + PROTOCOL_TCP, PROTOCOL_UDP + ), vol.Optional(CONF_PREFIX, default=DEFAULT_PREFIX): cv.string, } ) @@ -46,29 +53,34 @@ def setup(hass, config): host = conf.get(CONF_HOST) prefix = conf.get(CONF_PREFIX) port = conf.get(CONF_PORT) + protocol = conf.get(CONF_PROTOCOL) - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - try: - sock.connect((host, port)) - sock.shutdown(2) - _LOGGER.debug("Connection to Graphite possible") - except OSError: - _LOGGER.error("Not able to connect to Graphite") - return False + if protocol == PROTOCOL_TCP: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + sock.connect((host, port)) + sock.shutdown(2) + _LOGGER.debug("Connection to Graphite possible") + except OSError: + _LOGGER.error("Not able to connect to Graphite") + return False + else: + _LOGGER.debug("No connection check for UDP possible") - GraphiteFeeder(hass, host, port, prefix) + GraphiteFeeder(hass, host, port, protocol, prefix) return True class GraphiteFeeder(threading.Thread): """Feed data to Graphite.""" - def __init__(self, hass, host, port, prefix): + def __init__(self, hass, host, port, protocol, prefix): """Initialize the feeder.""" super().__init__(daemon=True) self._hass = hass self._host = host self._port = port + self._protocol = protocol # rstrip any trailing dots in case they think they need it self._prefix = prefix.rstrip(".") self._queue = queue.Queue() @@ -101,12 +113,16 @@ class GraphiteFeeder(threading.Thread): def _send_to_graphite(self, data): """Send data to Graphite.""" - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.settimeout(10) - sock.connect((self._host, self._port)) - sock.sendall(data.encode("ascii")) - sock.send(b"\n") - sock.close() + if self._protocol == PROTOCOL_TCP: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.settimeout(10) + sock.connect((self._host, self._port)) + sock.sendall(data.encode("ascii")) + sock.send(b"\n") + sock.close() + else: + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + sock.sendto(data.encode("ascii") + b"\n", (self._host, self._port)) def _report_attributes(self, entity_id, new_state): """Report the attributes.""" diff --git a/tests/components/graphite/test_init.py b/tests/components/graphite/test_init.py index 88be3723936..b7f5071813e 100644 --- a/tests/components/graphite/test_init.py +++ b/tests/components/graphite/test_init.py @@ -24,7 +24,7 @@ class TestGraphite(unittest.TestCase): def setup_method(self, method): """Set up things to be run when tests are started.""" self.hass = get_test_home_assistant() - self.gf = graphite.GraphiteFeeder(self.hass, "foo", 123, "ha") + self.gf = graphite.GraphiteFeeder(self.hass, "foo", 123, "tcp", "ha") def teardown_method(self, method): """Stop everything that was started.""" @@ -45,10 +45,23 @@ class TestGraphite(unittest.TestCase): assert setup_component(self.hass, graphite.DOMAIN, config) assert mock_gf.call_count == 1 - assert mock_gf.call_args == mock.call(self.hass, "foo", 123, "me") + assert mock_gf.call_args == mock.call(self.hass, "foo", 123, "tcp", "me") assert mock_socket.call_count == 1 assert mock_socket.call_args == mock.call(socket.AF_INET, socket.SOCK_STREAM) + @patch("socket.socket") + @patch("homeassistant.components.graphite.GraphiteFeeder") + def test_full_udp_config(self, mock_gf, mock_socket): + """Test setup with full configuration and UDP protocol.""" + config = { + "graphite": {"host": "foo", "port": 123, "protocol": "udp", "prefix": "me"} + } + + assert setup_component(self.hass, graphite.DOMAIN, config) + assert mock_gf.call_count == 1 + assert mock_gf.call_args == mock.call(self.hass, "foo", 123, "udp", "me") + assert mock_socket.call_count == 0 + @patch("socket.socket") @patch("homeassistant.components.graphite.GraphiteFeeder") def test_config_port(self, mock_gf, mock_socket): @@ -63,7 +76,7 @@ class TestGraphite(unittest.TestCase): def test_subscribe(self): """Test the subscription.""" fake_hass = mock.MagicMock() - gf = graphite.GraphiteFeeder(fake_hass, "foo", 123, "ha") + gf = graphite.GraphiteFeeder(fake_hass, "foo", 123, "tcp", "ha") fake_hass.bus.listen_once.has_calls( [ mock.call(EVENT_HOMEASSISTANT_START, gf.start_listen),