Update services.py

custom_components/pstryk/services.py

@@ -1,7 +1,6 @@
 """Services for Pstryk Energy integration."""
 import logging
 import voluptuous as vol
-import json
 from homeassistant.core import HomeAssistant, ServiceCall, callback
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.typing import ConfigType
@@ -11,6 +10,7 @@ from homeassistant.helpers.event import async_track_point_in_time
 from datetime import timedelta
 from homeassistant.util import dt as dt_util
 
+from .mqtt_common import publish_mqtt_prices, setup_periodic_mqtt_publish
 from .const import (
     DOMAIN, 
     DEFAULT_MQTT_TOPIC_BUY, 
@@ -79,46 +79,12 @@ async def async_setup_services(hass: HomeAssistant) -> None:
         
         # Publish for all entries or the specified one
         for entry in config_entries:
-            key = f"{entry.entry_id}_mqtt"
+            # Get topics from options or use override/default
-            mqtt_publisher = hass.data[DOMAIN].get(key)
+            mqtt_topic_buy = topic_buy_override or entry.options.get(CONF_MQTT_TOPIC_BUY, DEFAULT_MQTT_TOPIC_BUY)
+            mqtt_topic_sell = topic_sell_override or entry.options.get(CONF_MQTT_TOPIC_SELL, DEFAULT_MQTT_TOPIC_SELL)
             
-            # If publisher not found, create a temporary one
+            # Use common function to publish
-            if not mqtt_publisher:
+            success = await publish_mqtt_prices(hass, entry.entry_id, mqtt_topic_buy, mqtt_topic_sell)
-                # Check if we have buy coordinator
-                buy_coordinator = hass.data[DOMAIN].get(f"{entry.entry_id}_buy")
-                sell_coordinator = hass.data[DOMAIN].get(f"{entry.entry_id}_sell")
-                if not buy_coordinator or not sell_coordinator:
-                    _LOGGER.error("Price coordinators not found for entry %s", entry.entry_id)
-                    continue
-                    
-                # Dynamically import the publisher class
-                from .mqtt_publisher import PstrykMqttPublisher
-                
-                # Get topic from options or use override/default
-                mqtt_topic_buy = topic_buy_override
-                mqtt_topic_sell = topic_sell_override
-                if not mqtt_topic_buy:
-                    mqtt_topic_buy = entry.options.get(CONF_MQTT_TOPIC_BUY, DEFAULT_MQTT_TOPIC_BUY)
-                if not mqtt_topic_sell:
-                    mqtt_topic_sell = entry.options.get(CONF_MQTT_TOPIC_SELL, DEFAULT_MQTT_TOPIC_SELL)
-                    
-                # Create temporary publisher
-                mqtt_publisher = PstrykMqttPublisher(
-                    hass, 
-                    entry.entry_id, 
-                    mqtt_topic_buy,
-                    mqtt_topic_sell
-                )
-                await mqtt_publisher.async_initialize()
-            else:
-                # Use the override topics if provided
-                if topic_buy_override:
-                    mqtt_publisher.mqtt_topic_buy = topic_buy_override
-                if topic_sell_override:
-                    mqtt_publisher.mqtt_topic_sell = topic_sell_override
-                
-            # Publish prices
-            success = await mqtt_publisher.publish_prices()
             
             if success:
                 _LOGGER.info("Manual MQTT publish to EVCC completed for entry %s", entry.entry_id)
@@ -147,94 +113,25 @@ async def async_setup_services(hass: HomeAssistant) -> None:
                 
         # Process all entries or just the specified one
         for entry in config_entries:
-            # Get MQTT publisher
-            mqtt_publisher = hass.data[DOMAIN].get(f"{entry.entry_id}_mqtt")
-            
-            # Get buy and sell coordinators
-            buy_coordinator = hass.data[DOMAIN].get(f"{entry.entry_id}_buy")
-            sell_coordinator = hass.data[DOMAIN].get(f"{entry.entry_id}_sell")
-            
-            if not buy_coordinator or not buy_coordinator.data or not sell_coordinator or not sell_coordinator.data:
-                _LOGGER.error("No data available for entry %s", entry.entry_id)
-                continue
-            
-            # If no publisher, create a temporary one
-            if not mqtt_publisher:
-                # Import publisher class
-                from .mqtt_publisher import PstrykMqttPublisher
-                
             # Get topics from options or use overrides
             mqtt_topic_buy = topic_buy_override or entry.options.get(CONF_MQTT_TOPIC_BUY, DEFAULT_MQTT_TOPIC_BUY)
             mqtt_topic_sell = topic_sell_override or entry.options.get(CONF_MQTT_TOPIC_SELL, DEFAULT_MQTT_TOPIC_SELL)
             
-                # Create publisher
+            # First immediate publish
-                mqtt_publisher = PstrykMqttPublisher(
+            success = await publish_mqtt_prices(hass, entry.entry_id, mqtt_topic_buy, mqtt_topic_sell)
-                    hass, 
-                    entry.entry_id, 
-                    mqtt_topic_buy,
-                    mqtt_topic_sell
-                )
-                await mqtt_publisher.async_initialize()
-            else:
-                # Update topics if overrides provided
-                if topic_buy_override:
-                    mqtt_publisher.mqtt_topic_buy = topic_buy_override
-                if topic_sell_override:
-                    mqtt_publisher.mqtt_topic_sell = topic_sell_override
             
-            # Format prices
+            if not success:
-            buy_prices = mqtt_publisher._format_prices_for_evcc(buy_coordinator.data, "buy")
+                _LOGGER.error("Failed to publish initial retained messages for entry %s", entry.entry_id)
-            sell_prices = mqtt_publisher._format_prices_for_evcc(sell_coordinator.data, "sell")
-            
-            if not buy_prices or not sell_prices:
-                _LOGGER.error("No price data available to publish")
                 continue
             
-            # Sort prices
-            buy_prices.sort(key=lambda x: x["start"])
-            sell_prices.sort(key=lambda x: x["start"])
-            
-            # Create payloads
-            buy_payload = json.dumps(buy_prices)
-            sell_payload = json.dumps(sell_prices)
-            
-            # Get topics
-            buy_topic = mqtt_publisher.mqtt_topic_buy
-            sell_topic = mqtt_publisher.mqtt_topic_sell
-            
             _LOGGER.info(
-                "Setting up scheduled retain for buy topic %s and sell topic %s for %d hours", 
+                "Setting up scheduled retain for topics %s and %s for %d hours", 
-                buy_topic,
+                mqtt_topic_buy,
-                sell_topic,
+                mqtt_topic_sell,
                 retain_hours
             )
             
-            # First immediate publish
+            # Schedule periodic re-publishing
-            await hass.services.async_call(
-                "mqtt", 
-                "publish",
-                {
-                    "topic": buy_topic,
-                    "payload": buy_payload,
-                    "qos": 1,
-                    "retain": True
-                },
-                blocking=True
-            )
-            
-            await hass.services.async_call(
-                "mqtt", 
-                "publish",
-                {
-                    "topic": sell_topic,
-                    "payload": sell_payload,
-                    "qos": 1,
-                    "retain": True
-                },
-                blocking=True
-            )
-            
-            # Schedule periodic re-publishing to ensure message stays retained
             now = dt_util.now()
             end_time = now + timedelta(hours=retain_hours)
             
@@ -247,35 +144,14 @@ async def async_setup_services(hass: HomeAssistant) -> None:
             
             # Create a function that will republish and reschedule itself
             async def republish_retain(now=None):
-                # Publish with retain
+                """Republish retained messages periodically."""
-                await hass.services.async_call(
+                success = await publish_mqtt_prices(hass, entry.entry_id, mqtt_topic_buy, mqtt_topic_sell)
-                    "mqtt", 
-                    "publish",
-                    {
-                        "topic": buy_topic,
-                        "payload": buy_payload,
-                        "qos": 1,
-                        "retain": True
-                    },
-                    blocking=True
-                )
-                
-                await hass.services.async_call(
-                    "mqtt", 
-                    "publish",
-                    {
-                        "topic": sell_topic,
-                        "payload": sell_payload,
-                        "qos": 1,
-                        "retain": True
-                    },
-                    blocking=True
-                )
                 
+                if success:
                     current_time = dt_util.now()
                     _LOGGER.debug(
                         "Re-published retained messages (will continue until %s)",
-                    end_time.isoformat()
+                        end_time.strftime("%Y-%m-%d %H:%M:%S")
                     )
                     
                     # If we've reached the end time, stop