Add binary sensor to show whether clients are connected

Co-authored-by: Joshua Spence <josh@joshuaspence.com>

components/stream_server/__init__.py

@@ -26,7 +26,8 @@ DEPENDENCIES = ["uart", "network"]
 
 MULTI_CONF = True
 
-StreamServerComponent = cg.global_ns.class_("StreamServerComponent", cg.Component)
+ns = cg.global_ns
+StreamServerComponent = ns.class_("StreamServerComponent", cg.Component)
 
 
 def validate_buffer_size(buffer_size):

components/stream_server/binary_sensor.py

@@ -0,0 +1,43 @@
+# Copyright (C) 2023 Oxan van Leeuwen
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import esphome.codegen as cg
+import esphome.config_validation as cv
+from esphome.components import binary_sensor
+from esphome.const import (
+    DEVICE_CLASS_CONNECTIVITY,
+    ENTITY_CATEGORY_DIAGNOSTIC,
+)
+from . import ns, StreamServerComponent
+
+CONF_CONNECTED = "connected"
+CONF_STREAM_SERVER = "stream_server"
+
+CONFIG_SCHEMA = cv.Schema(
+    {
+        cv.GenerateID(CONF_STREAM_SERVER): cv.use_id(StreamServerComponent),
+        cv.Required(CONF_CONNECTED): binary_sensor.binary_sensor_schema(
+            device_class=DEVICE_CLASS_CONNECTIVITY,
+            entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        ),
+    }
+)
+
+
+async def to_code(config):
+    server = await cg.get_variable(config[CONF_STREAM_SERVER])
+
+    sens = await binary_sensor.new_binary_sensor(config[CONF_CONNECTED])
+    cg.add(server.set_connected_sensor(sens))

components/stream_server/sensor.py

@@ -0,0 +1,41 @@
+# Copyright (C) 2023 Oxan van Leeuwen
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import esphome.codegen as cg
+import esphome.config_validation as cv
+from esphome.components import sensor
+from esphome.const import (
+	ENTITY_CATEGORY_DIAGNOSTIC,
+)
+from . import ns, StreamServerComponent
+
+CONF_CONNECTION_COUNT = "connection_count"
+CONF_STREAM_SERVER = "stream_server"
+
+CONFIG_SCHEMA = cv.Schema(
+	{
+		cv.GenerateID(CONF_STREAM_SERVER): cv.use_id(StreamServerComponent),
+		cv.Required(CONF_CONNECTION_COUNT): sensor.sensor_schema(
+			accuracy_decimals=0,
+			entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+		)
+	}
+)
+
+async def to_code(config):
+	server = await cg.get_variable(config[CONF_STREAM_SERVER])
+
+	sens = await sensor.new_sensor(config[CONF_CONNECTION_COUNT])
+	cg.add(server.set_connection_count_sensor(sens))

components/stream_server/stream_server.cpp

@@ -53,6 +53,9 @@ void StreamServerComponent::loop() {
 void StreamServerComponent::dump_config() {
     ESP_LOGCONFIG(TAG, "Stream Server:");
     ESP_LOGCONFIG(TAG, "  Address: %s:%u", esphome::network::get_use_address().c_str(), this->port_);
+#ifdef USE_BINARY_SENSOR
+    LOG_BINARY_SENSOR("  ", "Connected:", this->connected_sensor_);
+#endif
 }
 
 void StreamServerComponent::on_shutdown() {
@@ -60,6 +63,13 @@ void StreamServerComponent::on_shutdown() {
         client.socket->shutdown(SHUT_RDWR);
 }
 
+void StreamServerComponent::publish_sensor() {
+#ifdef USE_BINARY_SENSOR
+    if (this->connected_sensor_)
+        this->connected_sensor_->publish_state(this->clients_.size() > 0);
+#endif
+}
+
 void StreamServerComponent::accept() {
     struct sockaddr_storage client_addr;
     socklen_t client_addrlen = sizeof(client_addr);
@@ -71,12 +81,16 @@ void StreamServerComponent::accept() {
     std::string identifier = socket->getpeername();
     this->clients_.emplace_back(std::move(socket), identifier, this->buf_head_);
     ESP_LOGD(TAG, "New client connected from %s", identifier.c_str());
+    this->publish_sensor();
 }
 
 void StreamServerComponent::cleanup() {
     auto discriminator = [](const Client &client) { return !client.disconnected; };
     auto last_client = std::partition(this->clients_.begin(), this->clients_.end(), discriminator);
-    this->clients_.erase(last_client, this->clients_.end());
+    if (last_client != this->clients_.end()) {
+        this->clients_.erase(last_client, this->clients_.end());
+        this->publish_sensor();
+    }
 }
 
 void StreamServerComponent::read() {

components/stream_server/stream_server.h

@@ -20,6 +20,10 @@
 #include "esphome/components/socket/socket.h"
 #include "esphome/components/uart/uart.h"
 
+#ifdef USE_BINARY_SENSOR
+#include "esphome/components/binary_sensor/binary_sensor.h"
+#endif
+
 #include <memory>
 #include <string>
 #include <vector>
@@ -31,6 +35,10 @@ public:
     void set_uart_parent(esphome::uart::UARTComponent *parent) { this->stream_ = parent; }
     void set_buffer_size(size_t size) { this->buf_size_ = size; }
 
+#ifdef USE_BINARY_SENSOR
+    void set_connected_sensor(esphome::binary_sensor::BinarySensor *connected) { this->connected_sensor_ = connected; }
+#endif
+
     void setup() override;
     void loop() override;
     void dump_config() override;
@@ -41,6 +49,8 @@ public:
     void set_port(uint16_t port) { this->port_ = port; }
 
 protected:
+    void publish_sensor();
+
     void accept();
     void cleanup();
     void read();
@@ -64,6 +74,10 @@ protected:
     uint16_t port_;
     size_t buf_size_;
 
+#ifdef USE_BINARY_SENSOR
+    esphome::binary_sensor::BinarySensor *connected_sensor_;
+#endif
+
     std::unique_ptr<uint8_t[]> buf_{};
     size_t buf_head_{0};
     size_t buf_tail_{0};