[tests] Use the shared press_and_wait and next_state in the action and light tests

esphome/components/online_image/download_buffer.cpp

@@ -1,6 +1,5 @@
 #include "download_buffer.h"
 #include "esphome/core/log.h"
-#include <algorithm>
 #include <cstring>
 
 namespace esphome::online_image {
@@ -35,46 +34,6 @@ size_t DownloadBuffer::read(size_t len) {
   return this->unread_;
 }
 
-void DownloadSource::compact() {
-  if (this->keep_read_ || this->pos_ == this->start_)
-    return;
-  this->buffer_.read(this->pos_ - this->start_);
-  this->start_ = this->pos_;
-}
-
-int32_t DownloadSource::read(uint8_t *buf, uint32_t len) {
-  const size_t offset = this->pos_ - this->start_;
-  const size_t n = std::min<size_t>(len, this->buffer_.unread() - offset);
-  memcpy(buf, this->buffer_.data(offset), n);
-  this->pos_ += n;
-  return static_cast<int32_t>(n);
-}
-
-bool DownloadSource::seek(uint32_t position) {
-  // Only what is still in the buffer can be read again; a decoder never skips ahead of the download
-  if (position < this->start_ || position > this->start_ + this->buffer_.unread())
-    return false;
-  this->pos_ = position;
-  return true;
-}
-
-bool DownloadBuffer::grow(size_t size) {
-  if (this->size_ >= size)
-    return true;
-  RAMAllocator<uint8_t> allocator;
-  uint8_t *buffer = allocator.allocate(size);
-  if (buffer == nullptr) {
-    ESP_LOGE(TAG, "allocation of %zu bytes failed. Biggest block in heap: %zu Bytes", size,
-             allocator.get_max_free_block_size());
-    return false;
-  }
-  memcpy(buffer, this->buffer_, this->unread_);
-  allocator.deallocate(this->buffer_, this->size_);
-  this->buffer_ = buffer;
-  this->size_ = size;
-  return true;
-}
-
 size_t DownloadBuffer::resize(size_t size) {
   if (this->size_ >= size) {
     // Avoid useless reallocations; if the buffer is big enough, don't reallocate.

esphome/components/online_image/download_buffer.h

@@ -1,6 +1,5 @@
 #pragma once
 
-#include "esphome/components/runtime_image/image_decoder.h"
 #include "esphome/core/helpers.h"
 #include <cstddef>
 #include <cstdint>
@@ -36,8 +35,6 @@ class DownloadBuffer {
 
   void reset() { this->unread_ = 0; }
   size_t resize(size_t size);
-  /// Make the buffer at least size bytes, keeping the unread data. Returns false if the allocation failed.
-  bool grow(size_t size);
 
  protected:
   uint8_t *buffer_;
@@ -46,35 +43,4 @@ class DownloadBuffer {
   size_t unread_;
 };
 
-/**
- * @brief Lets a decoder read a download as it arrives, from the download buffer.
- *
- * The decoder reads at its own position in the buffer. Bytes it has read stay there until compact() drops
- * them, and compact() keeps everything while set_keep_read(true) is in effect, so the decoder can seek back
- * while it reads the headers.
- */
-class DownloadSource : public runtime_image::DecodeSource {
- public:
-  explicit DownloadSource(DownloadBuffer &buffer) : buffer_(buffer) {}
-
-  /// Start again at the beginning of a new download, keeping what is read
-  void reset() {
-    this->start_ = 0;
-    this->pos_ = 0;
-    this->keep_read_ = true;
-  }
-  void set_keep_read(bool keep) { this->keep_read_ = keep; }
-  /// Drop the bytes the decoder has read, to make room for more of the download
-  void compact();
-
-  int32_t read(uint8_t *buf, uint32_t len) override;
-  bool seek(uint32_t position) override;
-
- protected:
-  DownloadBuffer &buffer_;
-  size_t start_{0};  // Position in the download of the first byte in the buffer
-  size_t pos_{0};    // Position in the download the decoder reads next
-  bool keep_read_{true};
-};
-
 }  // namespace esphome::online_image

esphome/components/online_image/image.py

@@ -22,8 +22,6 @@ DEPENDENCIES = ["http_request"]
 CODEOWNERS = ["@guillempages", "@clydebarrow"]
 
 CONF_ON_DOWNLOAD_FINISHED = "on_download_finished"
-CONF_ON_DOWNLOAD_STARTED = "on_download_started"
-CONF_ON_DOWNLOAD_PROGRESS = "on_download_progress"
 CONF_UPDATE = "update"
 
 online_image_ns = cg.esphome_ns.namespace("online_image")
@@ -45,9 +43,7 @@ ONLINE_IMAGE_SCHEMA = (
             cv.Optional(CONF_REQUEST_HEADERS): cv.All(
                 cv.Schema({cv.string: cv.templatable(cv.string)})
             ),
-            cv.Optional(CONF_ON_DOWNLOAD_STARTED): automation.validate_automation({}),
             cv.Optional(CONF_ON_DOWNLOAD_FINISHED): automation.validate_automation({}),
-            cv.Optional(CONF_ON_DOWNLOAD_PROGRESS): automation.validate_automation({}),
             cv.Optional(CONF_ON_ERROR): automation.validate_automation({}),
         }
     )
@@ -100,17 +96,9 @@ automation.register_apply_action(
 
 
 _CALLBACK_AUTOMATIONS = (
-    automation.CallbackAutomation(
-        CONF_ON_DOWNLOAD_STARTED, "add_on_download_started_callback"
-    ),
     automation.CallbackAutomation(
         CONF_ON_DOWNLOAD_FINISHED, "add_on_finished_callback", [(bool, "cached")]
     ),
-    automation.CallbackAutomation(
-        CONF_ON_DOWNLOAD_PROGRESS,
-        "add_on_progress_callback",
-        [(float, "progress"), (cg.size_t, "bytes_read")],
-    ),
     automation.CallbackAutomation(CONF_ON_ERROR, "add_on_error_callback"),
 )
 

esphome/components/online_image/online_image.cpp

@@ -4,7 +4,6 @@
 #include "esphome/core/helpers.h"
 #include "esphome/core/log.h"
 #include <algorithm>
-#include <cmath>
 #include <cstdio>
 
 ESPHOME_LOG_TAG(TAG, "online_image");
@@ -16,11 +15,6 @@ static const char *const IF_MODIFIED_SINCE_HEADER_NAME = "if-modified-since";
 
 namespace esphome::online_image {
 
-// Time one loop pass may spend reading, and decoding a JPEG as it downloads
-static constexpr uint32_t PASS_BUDGET_MS = 20;
-// Shortest time between two on_download_progress calls
-static constexpr uint32_t PROGRESS_INTERVAL_MS = 100;
-
 OnlineImage::OnlineImage(const std::string &url, int width, int height, runtime_image::ImageFormat format,
                          image::ImageType type, image::Transparency transparency, image::Image *placeholder,
                          uint32_t buffer_size, bool is_big_endian)
@@ -146,12 +140,8 @@ void OnlineImage::update() {
     return;
   }
 
-  // A JPEG of known size is decoded as it arrives, so the download buffer only has to hold its headers.
-  // Otherwise the decoder needs the complete image in the buffer.
-  this->streaming_ = total_size > 0 && total_size <= INT32_MAX && this->decoder_supports_source();
-  this->stream_started_ = false;
-  this->source_.reset();
-  if (!this->streaming_ && format == runtime_image::JPEG && total_size > this->download_buffer_.size() &&
+  // JPEG requires the complete image in the download buffer before decoding
+  if (format == runtime_image::JPEG && total_size > this->download_buffer_.size() &&
       this->download_buffer_.resize(total_size) == 0) {
     // resize() logged the failed allocation; the next update() allocates the configured buffer again
     this->end_connection_();
@@ -162,8 +152,6 @@ void OnlineImage::update() {
   ESP_LOGI(TAG, "Downloading image (Size: %zu)", total_size);
   this->start_time_ = millis();
   this->last_data_time_ = this->start_time_;
-  this->download_started_callback_.call();
-  this->report_progress_(true);
   this->enable_loop();
 }
 
@@ -181,11 +169,6 @@ void OnlineImage::loop() {
     return;
   }
 
-  if (this->streaming_) {
-    this->loop_stream_();
-    return;
-  }
-
   // Check if download is complete — use decoder's format-specific completion check
   // to handle both known content-length and chunked transfer encoding
   if (this->is_decode_finished() || (this->downloader_->content_length > 0 &&
@@ -196,22 +179,15 @@ void OnlineImage::loop() {
   }
 
   // Download and decode more data
-  if (this->download_buffer_.free_capacity() > 0) {
-    const uint32_t now = App.get_loop_component_start_time();
-    size_t received = 0;
-    int len;
-    do {
-      len = this->downloader_->read(this->download_buffer_.append(),
-                                    std::min(this->download_buffer_.free_capacity(), READ_CHUNK));
-      if (len <= 0)
-        break;
-      this->download_buffer_.write(len);
-      received += len;
-    } while (this->download_buffer_.free_capacity() != 0 && millis() - now < PASS_BUDGET_MS);
+  size_t available = this->download_buffer_.free_capacity();
+  if (available > 0) {
+    // Download in chunks to avoid blocking
+    available = std::min(available, this->download_buffer_initial_size_);
+    auto len = this->downloader_->read(this->download_buffer_.append(), available);
 
-    if (received > 0) {
-      this->last_data_time_ = now;
-      this->report_progress_(false);
+    if (len > 0) {
+      this->last_data_time_ = App.get_loop_component_start_time();
+      this->download_buffer_.write(len);
 
       // Feed data to decoder
       auto consumed = this->feed_data(this->download_buffer_.data(), this->download_buffer_.unread());
@@ -279,114 +255,6 @@ void OnlineImage::loop() {
   }
 }
 
-void OnlineImage::loop_stream_() {
-  const uint32_t now = App.get_loop_component_start_time();
-  const size_t total = this->downloader_->content_length;
-  if (this->stream_started_)
-    this->source_.compact();
-
-  if (this->downloader_->get_bytes_read() < total) {
-    if (this->download_buffer_.free_capacity() == 0) {
-      // compact() already ran, so a full buffer means the decoder read nothing in the last pass
-      if (now - this->last_data_time_ >= this->parent_->get_timeout()) {
-        ESP_LOGE(TAG, "Image does not fit in the %zu byte download buffer; increase buffer_size",
-                 this->download_buffer_.size());
-        this->fail_download_();
-        return;
-      }
-    } else {
-      int len;
-      do {
-        len = this->downloader_->read(this->download_buffer_.append(),
-                                      std::min(this->download_buffer_.free_capacity(), READ_CHUNK));
-        if (len <= 0)
-          break;
-        this->last_data_time_ = now;
-        this->download_buffer_.write(len);
-        this->report_progress_(false);
-      } while (this->downloader_->get_bytes_read() < total && this->download_buffer_.free_capacity() != 0 &&
-               millis() - now < PASS_BUDGET_MS);
-      if (len < 0) {
-        ESP_LOGE(TAG, "Error downloading image: %d", len);
-        this->fail_download_();
-        return;
-      } else if (now - this->last_data_time_ >= this->parent_->get_timeout()) {
-        ESP_LOGE(TAG, "No data received for %" PRIu32 " ms", this->parent_->get_timeout());
-        this->fail_download_();
-        return;
-      }
-    }
-  }
-  const bool arrived = this->downloader_->get_bytes_read() >= total;
-
-  if (!this->stream_started_) {
-    // The decoder reads the headers when it opens the image, so wait until they are in the buffer
-    if (!arrived && this->download_buffer_.free_capacity() != 0)
-      return;
-    const int err = this->begin_source(&this->source_);
-    if (err != 0) {
-      // Headers cut short by a full buffer fail to parse, or make the decoder seek past the buffer. Such an image
-      // still decodes once it is all in memory, as it did before JPEG images were decoded as they download.
-      if (!arrived && (err == runtime_image::DECODE_ERROR_INVALID_TYPE || err == runtime_image::DECODE_ERROR_READ)) {
-        ESP_LOGW(TAG,
-                 "The image headers do not fit the %zu byte download buffer, so the whole image is downloaded "
-                 "before it is decoded; increase buffer_size to decode it as it downloads",
-                 this->download_buffer_.size());
-        // grow() logs a failed allocation
-        if (this->download_buffer_.grow(total)) {
-          this->streaming_ = false;
-        } else {
-          this->fail_download_();
-        }
-        return;
-      }
-      ESP_LOGE(TAG, "Error decoding image: %s", runtime_image::decode_error_to_string(err));
-      this->fail_download_();
-      return;
-    }
-    this->source_.set_keep_read(false);
-    this->stream_started_ = true;
-    // Opening the image and allocating its buffer is enough for one pass
-    return;
-  }
-
-  do {
-    const int result = this->decode_step();
-    if (result == runtime_image::DECODE_STEP_MORE)
-      continue;
-    if (result == runtime_image::DECODE_STEP_DONE) {
-      this->finish_download_();
-    } else if (result == runtime_image::DECODE_STEP_NEED_DATA) {
-      if (!arrived)
-        return;
-      // JPEGDEC only waits for bytes short of the size it was given, so this guards against a hang
-      ESP_LOGE(TAG, "Download ended before the image was complete");
-      this->fail_download_();
-    } else {
-      ESP_LOGE(TAG, "Error decoding image: %s", runtime_image::decode_error_to_string(result));
-      this->fail_download_();
-    }
-    return;
-  } while (millis() - now < PASS_BUDGET_MS);
-}
-
-void OnlineImage::report_progress_(bool force) {
-  const uint32_t now = App.get_loop_component_start_time();
-  const size_t bytes = this->downloader_->get_bytes_read();
-  const size_t total = this->downloader_->content_length;
-  // The last bytes are always reported, so a progress bar reaches 100%
-  if (!force && now - this->last_progress_time_ < PROGRESS_INTERVAL_MS && (total == 0 || bytes < total))
-    return;
-  this->last_progress_time_ = now;
-  const float progress = total > 0 ? std::min(100.0f, bytes * 100.0f / total) : NAN;
-  this->download_progress_callback_.call(progress, bytes);
-}
-
-void OnlineImage::fail_download_() {
-  this->end_connection_();
-  this->download_error_callback_.call();
-}
-
 void OnlineImage::finish_download_() {
   // Must be read before end_decode() ends the decoding session
   bool decoder_finished = this->is_decode_finished();
@@ -420,8 +288,6 @@ void OnlineImage::end_connection_() {
     this->downloader_ = nullptr;
   }
   this->download_buffer_.reset();
-  this->streaming_ = false;
-  this->stream_started_ = false;
   this->disable_loop();
 }
 

esphome/components/online_image/online_image.h

@@ -16,10 +16,6 @@ using t_http_codes = enum {
   HTTP_CODE_NOT_FOUND = 404,
 };
 
-/// Largest download read(). A read() blocks until it has every byte it asked for, so this bounds how long one loop
-/// pass waits on a slow link.
-static constexpr size_t READ_CHUNK = 1460;
-
 /**
  * @brief Download an image from a given URL, and decode it using the specified decoder.
  * The image will then be stored in a buffer, so that it can be re-displayed without the
@@ -74,17 +70,9 @@ class OnlineImage final : public PollingComponent,
    */
   void release();
 
-  /// Called when the server has accepted the request and the image starts to download
-  template<typename F> void add_on_download_started_callback(F &&callback) {
-    this->download_started_callback_.add(std::forward<F>(callback));
-  }
   template<typename F> void add_on_finished_callback(F &&callback) {
     this->download_finished_callback_.add(std::forward<F>(callback));
   }
-  /// Called with the percentage downloaded (NAN when the server did not send the size) and the bytes read so far
-  template<typename F> void add_on_progress_callback(F &&callback) {
-    this->download_progress_callback_.add(std::forward<F>(callback));
-  }
   template<typename F> void add_on_error_callback(F &&callback) {
     this->download_error_callback_.add(std::forward<F>(callback));
   }
@@ -93,20 +81,12 @@ class OnlineImage final : public PollingComponent,
   bool validate_url_(const std::string &url);
   void end_connection_();
   void finish_download_();
-  /// loop() for a JPEG that is decoded as it downloads
-  void loop_stream_();
-  void fail_download_();
-  /// Reports download progress, at most every PROGRESS_INTERVAL_MS unless force is set
-  void report_progress_(bool force);
 
   CallbackManager<void(bool)> download_finished_callback_{};
   CallbackManager<void()> download_error_callback_{};
-  LazyCallbackManager<void()> download_started_callback_{};
-  LazyCallbackManager<void(float, size_t)> download_progress_callback_{};
 
   std::shared_ptr<http_request::HttpContainer> downloader_{nullptr};
   DownloadBuffer download_buffer_;
-  DownloadSource source_{download_buffer_};
   /**
    * This is the *initial* size of the download buffer, not the current size.
    * The download buffer can be resized at runtime; the download_buffer_initial_size_
@@ -129,9 +109,6 @@ class OnlineImage final : public PollingComponent,
 
   uint32_t start_time_{0};
   uint32_t last_data_time_{0};
-  uint32_t last_progress_time_{0};
-  bool streaming_{false};       // The decoder reads the download itself, as it arrives
-  bool stream_started_{false};  // The decoder has read the headers
 };
 
 }  // namespace esphome::online_image

esphome/components/runtime_image/image_decoder.cpp

@@ -2,6 +2,7 @@
 #include "runtime_image.h"
 #include "esphome/core/log.h"
 #include <algorithm>
+#include <cmath>
 
 namespace esphome::runtime_image {
 
@@ -10,8 +11,8 @@ ESPHOME_LOG_TAG(TAG, "image_decoder");
 bool ImageDecoder::set_size(int width, int height) {
   bool success = this->image_->resize(width, height) > 0;
   this->size_valid_ = success;
-  this->draw_width_ = width;
-  this->draw_height_ = height;
+  this->x_scale_ = static_cast<double>(this->image_->get_buffer_width()) / width;
+  this->y_scale_ = static_cast<double>(this->image_->get_buffer_height()) / height;
   return success;
 }
 
@@ -19,24 +20,10 @@ void ImageDecoder::draw(int x, int y, int w, int h, const Color &color) {
   if (!this->size_valid_) {
     return;
   }
-  // When shrinking, each target pixel is painted once, by the first source pixel that reaches it, rather than by every
-  // pixel that lands on it. Integer maths: the S3 has no double FPU.
-  const int buffer_width = this->image_->get_buffer_width();
-  const int buffer_height = this->image_->get_buffer_height();
-  int x1 = x;
-  int y1 = y;
-  int x2 = x + w;
-  int y2 = y + h;
-  if (buffer_width != this->draw_width_ || buffer_height != this->draw_height_) {
-    x1 = scale_boundary(x, this->draw_width_, buffer_width);
-    y1 = scale_boundary(y, this->draw_height_, buffer_height);
-    x2 = scale_boundary(x + w, this->draw_width_, buffer_width);
-    y2 = scale_boundary(y + h, this->draw_height_, buffer_height);
-  }
-  x2 = std::min(buffer_width, x2);
-  y2 = std::min(buffer_height, y2);
-  for (int i = x1; i < x2; i++) {
-    for (int j = y1; j < y2; j++) {
+  auto width = std::min(this->image_->get_buffer_width(), static_cast<int>(std::ceil((x + w) * this->x_scale_)));
+  auto height = std::min(this->image_->get_buffer_height(), static_cast<int>(std::ceil((y + h) * this->y_scale_)));
+  for (int i = x * this->x_scale_; i < width; i++) {
+    for (int j = y * this->y_scale_; j < height; j++) {
       this->image_->draw_pixel(i, j, color);
     }
   }

esphome/components/runtime_image/image_decoder.h

@@ -6,15 +6,6 @@
 
 namespace esphome::runtime_image {
 
-/// The first of `target` pixels that pixel n of `source` paints, rounded up: ceil(n * target / source). The spans of
-/// neighbouring pixels meet without overlapping, so a shrunk image paints each target pixel once.
-inline int scale_boundary(int n, int source, int target) {
-  if (source <= 0)
-    return 0;
-  return static_cast<int>((static_cast<uint32_t>(n) * static_cast<uint32_t>(target) + source - 1) /
-                          static_cast<uint32_t>(source));
-}
-
 enum DecodeError : int {
   DECODE_ERROR_INVALID_TYPE = -1,
   DECODE_ERROR_UNSUPPORTED_FORMAT = -2,
@@ -92,8 +83,8 @@ class ImageDecoder {
     this->expected_size_ = 0;
     this->decoded_bytes_ = 0;
     this->size_valid_ = true;
-    this->draw_width_ = 0;
-    this->draw_height_ = 0;
+    this->x_scale_ = 1.0;
+    this->y_scale_ = 1.0;
   }
 
   /**
@@ -193,9 +184,8 @@ class ImageDecoder {
   uint8_t *image_buffer_() const;
   bool image_is_big_endian_() const;
 
-  // Size of the image as the decoder draws it, which draw() spreads over the buffer
-  int draw_width_ = 0;
-  int draw_height_ = 0;
+  double x_scale_ = 1.0;
+  double y_scale_ = 1.0;
   RuntimeImage *image_;
   size_t expected_size_ = 0;  // Expected data size (0 if unknown)
   size_t decoded_bytes_ = 0;  // Bytes processed so far

esphome/components/runtime_image/jpeg_decoder.cpp

@@ -28,14 +28,16 @@ bool JpegDecoder::prepare_rgb565_scaling_(int width, int height) {
   this->source_height_ = height;
   this->x_boundaries_.clear();
   this->y_boundaries_.clear();
-  const int target_width = this->image_->get_buffer_width();
-  const int target_height = this->image_->get_buffer_height();
-  // The same boundaries as the generic draw() path, so an 8 to 4 scale keeps columns 0, 2, 4 and 6 in both
-  for (int x = 0; x <= width; x++) {
-    this->x_boundaries_.push_back(static_cast<uint16_t>(scale_boundary(x, width, target_width)));
+  const uint32_t target_width = this->image_->get_buffer_width();
+  const uint32_t target_height = this->image_->get_buffer_height();
+  // Boundary n is the first target pixel source pixel n covers, rounded up. When shrinking, the source
+  // pixel whose span lands on a target pixel paints it, so an 8 to 4 scale keeps columns 0, 2, 4 and 6;
+  // the generic draw() path rounds down and keeps the odd ones instead
+  for (uint32_t x = 0; x <= static_cast<uint32_t>(width); x++) {
+    this->x_boundaries_.push_back(static_cast<uint16_t>((x * target_width + width - 1) / width));
   }
-  for (int y = 0; y <= height; y++) {
-    this->y_boundaries_.push_back(static_cast<uint16_t>(scale_boundary(y, height, target_height)));
+  for (uint32_t y = 0; y <= static_cast<uint32_t>(height); y++) {
+    this->y_boundaries_.push_back(static_cast<uint16_t>((y * target_height + height - 1) / height));
   }
   return true;
 }
@@ -136,7 +138,7 @@ int HOT JpegDecoder::decode(uint8_t *buffer, size_t size) {
   const int err = this->start_();
   if (err != 0)
     return err;
-  if (!this->jpeg_.decode(0, 0, this->scale_option_()))
+  if (!this->jpeg_.decode(0, 0, 0))
     return this->fail_();
   this->decoded_bytes_ = size;
   this->jpeg_.close();
@@ -163,28 +165,10 @@ int JpegDecoder::start_() {
   } else {
     this->jpeg_.setPixelType(RGB8888);
   }
-  const int width = this->jpeg_.getWidth();
-  const int height = this->jpeg_.getHeight();
-  if (!this->set_size(width, height)) {
+  if (!this->set_size(this->jpeg_.getWidth(), this->jpeg_.getHeight())) {
     return DECODE_ERROR_OUT_OF_MEMORY;
   }
-  // JPEGDEC can decode at 1/2, 1/4 or 1/8 size for much less work. Use the smallest size that still has at least as
-  // many pixels as the buffer, so a large image shown small is not decoded in full only to be thrown away.
-  const int buffer_width = this->image_->get_buffer_width();
-  const int buffer_height = this->image_->get_buffer_height();
-  this->scale_ = 1;
-  while (this->scale_ < JPEG_SCALE_EIGHTH && width / (this->scale_ * 2) >= buffer_width &&
-         height / (this->scale_ * 2) >= buffer_height) {
-    this->scale_ *= 2;
-  }
-  const int decoded_width = (width + this->scale_ - 1) / this->scale_;
-  const int decoded_height = (height + this->scale_ - 1) / this->scale_;
-  this->draw_width_ = decoded_width;
-  this->draw_height_ = decoded_height;
-  if (this->scale_ != 1) {
-    ESP_LOGD(TAG, "Decoding at 1/%d size, %d x %d", this->scale_, decoded_width, decoded_height);
-  }
-  if (this->direct_rgb565_ && !this->prepare_rgb565_scaling_(decoded_width, decoded_height)) {
+  if (this->direct_rgb565_ && !this->prepare_rgb565_scaling_(this->jpeg_.getWidth(), this->jpeg_.getHeight())) {
     return DECODE_ERROR_OUT_OF_MEMORY;
   }
   return 0;
@@ -268,7 +252,7 @@ int JpegDecoder::begin_source(DecodeSource *source) {
     this->source_ = nullptr;
     return err;
   }
-  if (!this->jpeg_.decodeBegin(0, 0, this->scale_option_()))
+  if (!this->jpeg_.decodeBegin(0, 0, 0))
     return this->fail_();
   return 0;
 }

esphome/components/runtime_image/jpeg_decoder.h

@@ -31,16 +31,12 @@ class JpegDecoder : public ImageDecoder {
   int HOT decode_step() override;
 
   bool uses_direct_rgb565() const { return this->direct_rgb565_; }
-  /// JPEGDEC decodes the image at 1/get_scale() of its size: 1, 2, 4 or 8
-  int get_scale() const { return this->scale_; }
   void draw_rgb565(int x, int y, int w, int h, const uint8_t *pixels);
 
  protected:
   bool prepare_rgb565_scaling_(int width, int height);
   /// Set up an opened image for decoding. Returns 0, or a DecodeError.
   int start_();
-  /// The JPEGDEC option for scale_: JPEG_SCALE_HALF, _QUARTER and _EIGHTH are the divisors themselves
-  int scale_option_() const { return this->scale_ == 1 ? 0 : this->scale_; }
   /// The DecodeError for a failed read of the source or JPEGDEC's last error, closing the image
   int fail_();
 
@@ -53,7 +49,6 @@ class JpegDecoder : public ImageDecoder {
   FixedVector<uint16_t> y_boundaries_{};
   int source_width_{0};
   int source_height_{0};
-  int scale_{1};                   // JPEGDEC decodes at 1/scale_ of the image size
   DecodeSource *source_{nullptr};  // Set while a begin_source() decode is in progress
   bool direct_rgb565_{false};
   bool source_failed_{false};

esphome/components/runtime_image/runtime_image.cpp

@@ -116,9 +116,7 @@ void RuntimeImage::draw_pixel(int x, int y, const Color &color) {
     }
     case image::IMAGE_TYPE_GRAYSCALE: {
       uint32_t pos = this->get_position_(x, y);
-      // 0.2125, 0.7154 and 0.0721 in 1/256ths, summing to 256 so white stays 255. Integer maths: the S3 has no
-      // double FPU
-      auto gray = static_cast<uint8_t>((54 * color.r + 183 * color.g + 19 * color.b) >> 8);
+      auto gray = static_cast<uint8_t>(0.2125 * color.r + 0.7154 * color.g + 0.0721 * color.b);
       if (this->transparency_ == image::TRANSPARENCY_CHROMA_KEY) {
         if (gray == 1) {
           gray = 0;

tests/components/online_image/common.yaml

@@ -19,11 +19,6 @@ image:
         } else {
           ESP_LOGD("online_image", "Cache miss: fresh download");
         }
-    on_download_started:
-      - logger.log: Download started
-    on_download_progress:
-      lambda: |-
-        ESP_LOGD("online_image", "Downloaded %.0f%% (%zu bytes)", progress, bytes_read);
   - platform: online_image
     id: online_binary_transparent_image
     url: http://www.libpng.org/pub/png/img_png/pnglogo-blk-tiny.png

tests/components/online_image/test_download_end.cpp

@@ -2,7 +2,6 @@
 
 #include <algorithm>
 #include <array>
-#include <cmath>
 #include <cstdint>
 #include <cstring>
 #include <memory>
@@ -28,9 +27,6 @@ class FakeContainer : public http_request::HttpContainer {
   }
 
   int read(uint8_t *buf, size_t max_len) override {
-    this->largest_read = std::max(this->largest_read, max_len);
-    if (this->read_error_)
-      return -1;
     size_t len = std::min(max_len, this->available_ - this->bytes_read_);
     std::memcpy(buf, this->body_.data() + this->bytes_read_, len);
     this->bytes_read_ += len;
@@ -46,21 +42,13 @@ class FakeContainer : public http_request::HttpContainer {
   }
 
   void set_content_length(size_t length) { this->content_length = length; }
-  /// Make every read from now on fail, like a dropped connection
-  void fail_reads() { this->read_error_ = true; }
-  /// Let len more bytes of the body arrive
-  void deliver(size_t len) { this->available_ = std::min(this->body_.size(), this->available_ + len); }
   void add_response_header(const std::string &name, const std::string &value) {
     this->response_headers_.push_back({name, value});
   }
 
-  /// The most bytes any one read() asked for
-  size_t largest_read{0};
-
  protected:
   std::vector<uint8_t> body_;
   size_t available_;
-  bool read_error_{false};
 };
 
 class FakeHttpRequest : public http_request::HttpRequestComponent {
@@ -135,173 +123,13 @@ static constexpr std::array<uint8_t, 69> PNG_1X1{
     0x0C, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9C, 0x63, 0xF8, 0xCF, 0xC0, 0x00, 0x00, 0x03, 0x01, 0x01, 0x00, 0xC9,
     0xFE, 0x92, 0xEF, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82};
 
-// 8x8 progressive JPEG, 521 bytes
-static const uint8_t JPEG_PROGRESSIVE[] = {
-    0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10, 0x4A, 0x46, 0x49, 0x46, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00,
-    0x00, 0xFF, 0xDB, 0x00, 0x43, 0x00, 0x10, 0x0B, 0x0C, 0x0E, 0x0C, 0x0A, 0x10, 0x0E, 0x0D, 0x0E, 0x12, 0x11, 0x10,
-    0x13, 0x18, 0x28, 0x1A, 0x18, 0x16, 0x16, 0x18, 0x31, 0x23, 0x25, 0x1D, 0x28, 0x3A, 0x33, 0x3D, 0x3C, 0x39, 0x33,
-    0x38, 0x37, 0x40, 0x48, 0x5C, 0x4E, 0x40, 0x44, 0x57, 0x45, 0x37, 0x38, 0x50, 0x6D, 0x51, 0x57, 0x5F, 0x62, 0x67,
-    0x68, 0x67, 0x3E, 0x4D, 0x71, 0x79, 0x70, 0x64, 0x78, 0x5C, 0x65, 0x67, 0x63, 0xFF, 0xDB, 0x00, 0x43, 0x01, 0x11,
-    0x12, 0x12, 0x18, 0x15, 0x18, 0x2F, 0x1A, 0x1A, 0x2F, 0x63, 0x42, 0x38, 0x42, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63,
-    0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63,
-    0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63,
-    0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0xFF, 0xC2, 0x00, 0x11, 0x08, 0x00, 0x08, 0x00, 0x08, 0x03, 0x01, 0x22, 0x00,
-    0x02, 0x11, 0x01, 0x03, 0x11, 0x01, 0xFF, 0xC4, 0x00, 0x15, 0x00, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xFF, 0xC4, 0x00, 0x15, 0x01, 0x01, 0x01, 0x00, 0x00,
-    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0x05, 0xFF, 0xDA, 0x00, 0x0C, 0x03,
-    0x01, 0x00, 0x02, 0x10, 0x03, 0x10, 0x00, 0x00, 0x01, 0x94, 0x05, 0x3F, 0xFF, 0xC4, 0x00, 0x14, 0x10, 0x01, 0x00,
-    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xDA, 0x00, 0x08,
-    0x01, 0x01, 0x00, 0x01, 0x05, 0x02, 0x7F, 0xFF, 0xC4, 0x00, 0x14, 0x11, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xDA, 0x00, 0x08, 0x01, 0x03, 0x01, 0x01, 0x3F,
-    0x01, 0x7F, 0xFF, 0xC4, 0x00, 0x14, 0x11, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-    0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xDA, 0x00, 0x08, 0x01, 0x02, 0x01, 0x01, 0x3F, 0x01, 0x7F, 0xFF, 0xC4, 0x00,
-    0x14, 0x10, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-    0xFF, 0xDA, 0x00, 0x08, 0x01, 0x01, 0x00, 0x06, 0x3F, 0x02, 0x7F, 0xFF, 0xC4, 0x00, 0x14, 0x10, 0x01, 0x00, 0x00,
-    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xDA, 0x00, 0x08, 0x01,
-    0x01, 0x00, 0x01, 0x3F, 0x21, 0x7F, 0xFF, 0xDA, 0x00, 0x0C, 0x03, 0x01, 0x00, 0x02, 0x00, 0x03, 0x00, 0x00, 0x00,
-    0x10, 0x0B, 0xFF, 0xC4, 0x00, 0x14, 0x11, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-    0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xDA, 0x00, 0x08, 0x01, 0x03, 0x01, 0x01, 0x3F, 0x10, 0x7F, 0xFF, 0xC4, 0x00,
-    0x14, 0x11, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-    0xFF, 0xDA, 0x00, 0x08, 0x01, 0x02, 0x01, 0x01, 0x3F, 0x10, 0x7F, 0xFF, 0xC4, 0x00, 0x14, 0x10, 0x01, 0x00, 0x00,
-    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xDA, 0x00, 0x08, 0x01,
-    0x01, 0x00, 0x01, 0x3F, 0x10, 0x7F, 0xFF, 0xD9,
-};
-
-// 48x40 4:2:0 JPEG (quality 95), 2344 bytes: three rows of 16 pixel MCUs, with the scan starting at byte 609
-static const uint8_t JPEG_48X40[] = {
-    0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10, 0x4A, 0x46, 0x49, 0x46, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00,
-    0x00, 0xFF, 0xDB, 0x00, 0x43, 0x00, 0x02, 0x01, 0x01, 0x01, 0x01, 0x01, 0x02, 0x01, 0x01, 0x01, 0x02, 0x02, 0x02,
-    0x02, 0x02, 0x04, 0x03, 0x02, 0x02, 0x02, 0x02, 0x05, 0x04, 0x04, 0x03, 0x04, 0x06, 0x05, 0x06, 0x06, 0x06, 0x05,
-    0x06, 0x06, 0x06, 0x07, 0x09, 0x08, 0x06, 0x07, 0x09, 0x07, 0x06, 0x06, 0x08, 0x0B, 0x08, 0x09, 0x0A, 0x0A, 0x0A,
-    0x0A, 0x0A, 0x06, 0x08, 0x0B, 0x0C, 0x0B, 0x0A, 0x0C, 0x09, 0x0A, 0x0A, 0x0A, 0xFF, 0xDB, 0x00, 0x43, 0x01, 0x02,
-    0x02, 0x02, 0x02, 0x02, 0x02, 0x05, 0x03, 0x03, 0x05, 0x0A, 0x07, 0x06, 0x07, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A,
-    0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A,
-    0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A,
-    0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0x0A, 0xFF, 0xC0, 0x00, 0x11, 0x08, 0x00, 0x28, 0x00, 0x30, 0x03, 0x01, 0x22, 0x00,
-    0x02, 0x11, 0x01, 0x03, 0x11, 0x01, 0xFF, 0xC4, 0x00, 0x1F, 0x00, 0x00, 0x01, 0x05, 0x01, 0x01, 0x01, 0x01, 0x01,
-    0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A,
-    0x0B, 0xFF, 0xC4, 0x00, 0xB5, 0x10, 0x00, 0x02, 0x01, 0x03, 0x03, 0x02, 0x04, 0x03, 0x05, 0x05, 0x04, 0x04, 0x00,
-    0x00, 0x01, 0x7D, 0x01, 0x02, 0x03, 0x00, 0x04, 0x11, 0x05, 0x12, 0x21, 0x31, 0x41, 0x06, 0x13, 0x51, 0x61, 0x07,
-    0x22, 0x71, 0x14, 0x32, 0x81, 0x91, 0xA1, 0x08, 0x23, 0x42, 0xB1, 0xC1, 0x15, 0x52, 0xD1, 0xF0, 0x24, 0x33, 0x62,
-    0x72, 0x82, 0x09, 0x0A, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A, 0x34, 0x35, 0x36, 0x37,
-    0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49, 0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A,
-    0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x83, 0x84, 0x85,
-    0x86, 0x87, 0x88, 0x89, 0x8A, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6,
-    0xA7, 0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7,
-    0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA, 0xE1, 0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7,
-    0xE8, 0xE9, 0xEA, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8, 0xF9, 0xFA, 0xFF, 0xC4, 0x00, 0x1F, 0x01, 0x00,
-    0x03, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x02, 0x03,
-    0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0xFF, 0xC4, 0x00, 0xB5, 0x11, 0x00, 0x02, 0x01, 0x02, 0x04, 0x04,
-    0x03, 0x04, 0x07, 0x05, 0x04, 0x04, 0x00, 0x01, 0x02, 0x77, 0x00, 0x01, 0x02, 0x03, 0x11, 0x04, 0x05, 0x21, 0x31,
-    0x06, 0x12, 0x41, 0x51, 0x07, 0x61, 0x71, 0x13, 0x22, 0x32, 0x81, 0x08, 0x14, 0x42, 0x91, 0xA1, 0xB1, 0xC1, 0x09,
-    0x23, 0x33, 0x52, 0xF0, 0x15, 0x62, 0x72, 0xD1, 0x0A, 0x16, 0x24, 0x34, 0xE1, 0x25, 0xF1, 0x17, 0x18, 0x19, 0x1A,
-    0x26, 0x27, 0x28, 0x29, 0x2A, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49, 0x4A,
-    0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A, 0x73, 0x74, 0x75,
-    0x76, 0x77, 0x78, 0x79, 0x7A, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x92, 0x93, 0x94, 0x95, 0x96,
-    0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7, 0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7,
-    0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8,
-    0xD9, 0xDA, 0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8, 0xF9,
-    0xFA, 0xFF, 0xDA, 0x00, 0x0C, 0x03, 0x01, 0x00, 0x02, 0x11, 0x03, 0x11, 0x00, 0x3F, 0x00, 0xFC, 0xAA, 0xB8, 0xB7,
-    0xD3, 0xB5, 0xD8, 0x1E, 0xD6, 0xD1, 0xA5, 0x5B, 0x4B, 0x46, 0x10, 0xC1, 0x06, 0xA1, 0x29, 0x69, 0x25, 0xB7, 0xC2,
-    0xC6, 0x0C, 0x71, 0x26, 0x5E, 0x50, 0x15, 0x5D, 0xB6, 0x9C, 0x9E, 0x00, 0xC1, 0x67, 0x67, 0x4D, 0xCF, 0x84, 0xF7,
-    0x57, 0xA6, 0xDA, 0x08, 0x52, 0xF5, 0xE5, 0xBA, 0xB8, 0x7B, 0x5B, 0x68, 0xEE, 0x2F, 0x66, 0x89, 0xB7, 0xB7, 0x9B,
-    0xE5, 0x9C, 0x23, 0xAA, 0xFE, 0xF7, 0x68, 0x11, 0x9E, 0x5D, 0x18, 0x12, 0x48, 0x56, 0x5C, 0x9F, 0x57, 0xB6, 0xF8,
-    0x33, 0xF0, 0xE2, 0x5B, 0xFB, 0xCF, 0xF4, 0x51, 0x0C, 0xC6, 0x64, 0x30, 0x49, 0x71, 0x7C, 0xE2, 0x16, 0x90, 0x2E,
-    0xF4, 0x72, 0xEA, 0xE0, 0x34, 0x81, 0x77, 0x93, 0x9C, 0x07, 0xC2, 0x93, 0xB9, 0x61, 0xC0, 0xD3, 0xF0, 0x17, 0xC1,
-    0xBF, 0x03, 0xCD, 0xE1, 0xAD, 0x42, 0x3F, 0x0A, 0xD8, 0xC7, 0x2D, 0xAC, 0x36, 0xCB, 0xF6, 0x79, 0x22, 0x69, 0x84,
-    0xC4, 0x93, 0x29, 0x2A, 0x89, 0xB9, 0x00, 0x50, 0x62, 0x7C, 0xFF, 0x00, 0x16, 0xD9, 0x65, 0x00, 0x16, 0x65, 0x07,
-    0xFA, 0x5B, 0x36, 0xE2, 0xAC, 0xBF, 0x2D, 0x55, 0x79, 0x29, 0xB5, 0x26, 0x9B, 0xA9, 0x18, 0xB5, 0x15, 0x36, 0x9B,
-    0x53, 0xD1, 0xCA, 0x10, 0x93, 0x69, 0xB6, 0x94, 0x27, 0x16, 0xAF, 0x19, 0x29, 0xB8, 0xBB, 0x3F, 0x9A, 0xE1, 0x1E,
-    0x2A, 0xCB, 0x9E, 0x79, 0x4A, 0x92, 0x84, 0xDA, 0xD6, 0xED, 0xA4, 0xED, 0x7A, 0x72, 0xB5, 0xAC, 0xF6, 0xEB, 0xAD,
-    0xD5, 0x97, 0x91, 0xC9, 0x4D, 0xA4, 0xE9, 0xD7, 0x56, 0x45, 0x66, 0xD3, 0x9D, 0x85, 0xB1, 0x86, 0x59, 0x7E, 0xD9,
-    0x7E, 0xA0, 0xE0, 0xA8, 0x02, 0x17, 0x02, 0x3C, 0x36, 0x59, 0x18, 0xE2, 0x42, 0x1D, 0x40, 0x56, 0x0A, 0x77, 0x30,
-    0xAE, 0xFA, 0x56, 0xD4, 0xA4, 0x8A, 0xD9, 0xAD, 0xD6, 0x38, 0x1D, 0x7F, 0x77, 0x77, 0x1C, 0xB2, 0x4C, 0x22, 0x8C,
-    0x2B, 0xF3, 0x14, 0x64, 0x45, 0x85, 0x52, 0x47, 0x92, 0x7E, 0xF2, 0x0F, 0x34, 0x1D, 0xBC, 0x56, 0xDE, 0x85, 0xF0,
-    0x4B, 0xE1, 0x0A, 0xE9, 0xB0, 0xDB, 0x2D, 0x8C, 0x4C, 0xCD, 0x74, 0xF1, 0x42, 0x96, 0x93, 0xCF, 0xF6, 0x54, 0x0C,
-    0xB3, 0x44, 0x03, 0xAA, 0x6C, 0x49, 0x10, 0x87, 0x00, 0x88, 0xC2, 0xEE, 0x66, 0x6F, 0x9B, 0xAB, 0x57, 0xA4, 0x68,
-    0xDF, 0x0A, 0x3E, 0x1B, 0xDA, 0xE9, 0xDF, 0xD9, 0x96, 0xFE, 0x03, 0xB2, 0x8E, 0xD9, 0x8A, 0xDD, 0xC1, 0x77, 0x14,
-    0x97, 0x50, 0x2E, 0xE3, 0x20, 0x62, 0xA0, 0x33, 0x0C, 0xC6, 0x58, 0x48, 0x36, 0x6D, 0xC1, 0x20, 0xA2, 0x2B, 0x00,
-    0xB5, 0xF1, 0x99, 0xCF, 0x19, 0x65, 0x79, 0xAD, 0x0C, 0x3C, 0x2B, 0x61, 0xF9, 0x28, 0xD6, 0x52, 0x97, 0x35, 0x95,
-    0xD4, 0x61, 0x28, 0x5B, 0x99, 0xB8, 0xA9, 0xA9, 0x47, 0xD9, 0x3A, 0x69, 0xA9, 0xC6, 0x51, 0x52, 0x6A, 0x0D, 0x4A,
-    0x74, 0xA9, 0xC3, 0xFA, 0x7F, 0xC3, 0xFE, 0x25, 0xC1, 0x4E, 0xAE, 0x26, 0x2E, 0x13, 0x71, 0x5C, 0x8D, 0xFB, 0xB1,
-    0xD6, 0xEA, 0x6F, 0xF9, 0xAF, 0xBE, 0xB7, 0x4B, 0x55, 0x1D, 0x75, 0x4C, 0xF3, 0xDF, 0x05, 0xD9, 0xF8, 0x7B, 0xC3,
-    0xBA, 0xB5, 0xF1, 0xBC, 0x8D, 0x63, 0xDF, 0x1C, 0x32, 0x88, 0x1A, 0xF3, 0x74, 0x73, 0x89, 0x15, 0xC4, 0x8A, 0x25,
-    0x66, 0xDC, 0xAA, 0x98, 0x60, 0x06, 0x7F, 0xBD, 0xF2, 0x81, 0xF3, 0x36, 0xFE, 0x9D, 0xAF, 0x78, 0x8F, 0x54, 0xD3,
-    0xAD, 0xA6, 0x9E, 0xC5, 0xE5, 0x9C, 0x96, 0xDC, 0xB7, 0xE8, 0x84, 0x47, 0xFB, 0xB4, 0x54, 0x42, 0x0F, 0x03, 0xF8,
-    0xD7, 0xE5, 0x39, 0x60, 0xF8, 0x04, 0x79, 0x7B, 0xE3, 0xF4, 0x9F, 0x85, 0x5F, 0x08, 0xBE, 0x0D, 0x2C, 0x57, 0xCE,
-    0x6C, 0xA5, 0x2B, 0xA8, 0x08, 0x20, 0x82, 0x04, 0xB9, 0xB8, 0x82, 0x43, 0x28, 0x72, 0x55, 0xD5, 0x83, 0xB9, 0x05,
-    0xC2, 0x44, 0x55, 0x77, 0x8C, 0xFC, 0xEC, 0xD8, 0xCB, 0xE7, 0xBE, 0x9B, 0xE0, 0x5F, 0xC3, 0x9D, 0x32, 0xD4, 0x45,
-    0x26, 0x8B, 0x2C, 0x93, 0xC0, 0xF3, 0x95, 0x49, 0xEE, 0x2E, 0x0C, 0x26, 0x26, 0x31, 0xC9, 0xF3, 0xED, 0x95, 0x96,
-    0x32, 0xEC, 0x10, 0xE1, 0xB8, 0x60, 0x3A, 0x90, 0x1B, 0x3F, 0x9E, 0x66, 0x58, 0xFC, 0x14, 0x31, 0x15, 0xFD, 0x9C,
-    0x79, 0x57, 0x3A, 0x92, 0x84, 0x1C, 0x3D, 0xC9, 0xCD, 0x27, 0x6D, 0x12, 0xE5, 0xE7, 0xA6, 0xD5, 0x4D, 0x64, 0x93,
-    0xD1, 0x45, 0xC6, 0xAD, 0xE7, 0x2F, 0xE9, 0xEE, 0x17, 0xE3, 0x0C, 0x0A, 0xE1, 0x9C, 0x37, 0x34, 0x67, 0x2B, 0x73,
-    0x2B, 0xB4, 0xBA, 0x49, 0xF4, 0x72, 0x57, 0x51, 0xDD, 0x24, 0x9B, 0x5A, 0x36, 0x9A, 0x49, 0x1E, 0x15, 0xA7, 0x7E,
-    0xCE, 0xF6, 0x89, 0x6A, 0x96, 0x43, 0x55, 0xB5, 0xB3, 0xB6, 0xB8, 0x96, 0x2B, 0x9B, 0x78, 0xA0, 0xB4, 0x9E, 0xE5,
-    0x61, 0x90, 0xE0, 0xC2, 0xEC, 0x61, 0x91, 0x47, 0xCC, 0xDB, 0xF1, 0x80, 0xC0, 0x98, 0xD4, 0xE5, 0xDC, 0x64, 0xEC,
-    0x7C, 0x29, 0xFD, 0x98, 0x6C, 0xAF, 0xEC, 0xF5, 0x76, 0xB7, 0xF1, 0x36, 0x9A, 0x5D, 0x6F, 0x96, 0xE1, 0x74, 0xB5,
-    0xD3, 0x44, 0x91, 0xA8, 0x45, 0x0C, 0xE2, 0x3D, 0xAC, 0x85, 0xF9, 0x3B, 0x5B, 0x19, 0x8A, 0x43, 0x1F, 0xCA, 0x50,
-    0x9D, 0x87, 0xD2, 0x0F, 0xC2, 0x1F, 0x89, 0x4F, 0x6D, 0x04, 0x0F, 0xE1, 0x61, 0xAA, 0xEE, 0xB4, 0x8E, 0x56, 0xFB,
-    0x06, 0xAF, 0x04, 0xF3, 0x40, 0x9B, 0x55, 0x4E, 0xCB, 0x66, 0x91, 0xCB, 0x18, 0xC2, 0x88, 0xB6, 0xB1, 0x65, 0x60,
-    0xA4, 0x6E, 0xDC, 0xA4, 0x9E, 0xAF, 0xE1, 0xBF, 0xC2, 0x2F, 0x89, 0xA9, 0xE1, 0xFB, 0x8D, 0x6E, 0xC2, 0x3B, 0x4B,
-    0x30, 0xD7, 0x60, 0xDA, 0x41, 0x2D, 0xCC, 0x6C, 0x2E, 0x49, 0xF3, 0x1E, 0x42, 0x19, 0x5D, 0x58, 0x3C, 0x8A, 0x81,
-    0x82, 0xB8, 0x0C, 0xC9, 0xF2, 0xE5, 0x8F, 0x03, 0x8F, 0x36, 0xE3, 0x8A, 0x98, 0xCC, 0x14, 0xAB, 0x62, 0x2B, 0xC2,
-    0x34, 0xA2, 0xA5, 0x78, 0x37, 0x42, 0x12, 0x70, 0x8C, 0x62, 0xD2, 0x6D, 0x72, 0x7B, 0xAF, 0x9B, 0xDD, 0xFD, 0xE5,
-    0x4E, 0x79, 0x49, 0x45, 0x73, 0xA9, 0x72, 0x2F, 0xF0, 0x73, 0x83, 0xF8, 0xCE, 0xB4, 0x33, 0xE8, 0x3F, 0xAE, 0x42,
-    0xF7, 0x76, 0xB3, 0x83, 0x57, 0xE5, 0x95, 0xBE, 0xCE, 0x8F, 0x7B, 0xDE, 0xD7, 0x4B, 0x6B, 0x9C, 0x1E, 0x8D, 0xF0,
-    0x56, 0x7B, 0xDD, 0x26, 0x34, 0x5F, 0x1F, 0xDB, 0x9D, 0x37, 0x53, 0xB8, 0x96, 0xED, 0xA0, 0x16, 0xB1, 0xC9, 0x2C,
-    0x8B, 0xF2, 0x26, 0x58, 0x89, 0x37, 0xEE, 0x04, 0xB2, 0x97, 0x54, 0x2D, 0x88, 0xC0, 0x60, 0x0B, 0x03, 0x5E, 0x9B,
-    0x67, 0xFB, 0x36, 0xE9, 0x31, 0x2D, 0x9C, 0x37, 0x9E, 0x32, 0x84, 0x5D, 0x2C, 0x4A, 0xD6, 0xAD, 0x73, 0x63, 0x28,
-    0x36, 0xAA, 0x53, 0xF7, 0x90, 0x29, 0x8E, 0x5C, 0x92, 0xBE, 0x41, 0xDC, 0xC0, 0xB3, 0x2B, 0x46, 0x32, 0x4B, 0x96,
-    0x35, 0xAD, 0x1F, 0xC2, 0xDF, 0x89, 0x17, 0xF6, 0xB6, 0xFA, 0xB5, 0xED, 0xB8, 0xB7, 0x8E, 0x78, 0x66, 0x96, 0x09,
-    0x6C, 0x6E, 0x23, 0x32, 0x28, 0x08, 0xE8, 0xF2, 0xBA, 0xF9, 0x86, 0x48, 0xC0, 0x68, 0xFE, 0x63, 0xC8, 0x23, 0x82,
-    0xCC, 0xA5, 0x4A, 0xFA, 0x66, 0x8F, 0xF0, 0x37, 0xE2, 0x6D, 0xE5, 0xC9, 0x3A, 0x87, 0x86, 0xB4, 0xFB, 0x28, 0xAE,
-    0x61, 0x0F, 0x0B, 0x59, 0xDF, 0xC0, 0xD1, 0xAC, 0xE0, 0x24, 0x41, 0x42, 0xC4, 0x49, 0x1B, 0x10, 0x21, 0x50, 0xC0,
-    0x09, 0x11, 0x8B, 0xEE, 0x6E, 0x71, 0xE2, 0xE2, 0x78, 0xBB, 0x33, 0xFA, 0xC5, 0x1A, 0xB4, 0xEA, 0xC6, 0x74, 0xD2,
-    0x8C, 0x24, 0xA5, 0x2A, 0x3C, 0xB5, 0x2E, 0xB9, 0x9B, 0x71, 0x97, 0x34, 0xE0, 0xAA, 0x54, 0x93, 0x71, 0x56, 0x49,
-    0xB9, 0xC6, 0x4E, 0x9A, 0x9B, 0xD3, 0xFA, 0x7F, 0x81, 0xB8, 0xD6, 0xAC, 0xDD, 0x6F, 0x65, 0x8D, 0x82, 0x77, 0x85,
-    0x95, 0xE1, 0x7D, 0x1C, 0xD2, 0x8B, 0x7A, 0x27, 0xA2, 0xF4, 0xED, 0xEE, 0x9E, 0x67, 0xE0, 0x8F, 0x81, 0x37, 0xBA,
-    0xEF, 0x87, 0x67, 0xBF, 0x8F, 0xC4, 0xD2, 0xDD, 0x4F, 0x05, 0xA8, 0x61, 0x7B, 0x1E, 0x9E, 0xE1, 0x8C, 0xEA, 0xF3,
-    0x18, 0x99, 0xA3, 0xDE, 0x76, 0x66, 0x36, 0x2A, 0x30, 0xFB, 0x9B, 0x0F, 0xC3, 0x1C, 0xA3, 0x76, 0x4B, 0xFB, 0x3B,
-    0xE8, 0xBA, 0x66, 0xBF, 0x0E, 0x93, 0xFD, 0xAB, 0x23, 0xD9, 0xA5, 0xCE, 0x62, 0xF3, 0x34, 0xEB, 0x7F, 0x29, 0x0A,
-    0xA7, 0x04, 0x48, 0xF2, 0x79, 0x82, 0x42, 0x1D, 0x46, 0x4B, 0xA9, 0x1B, 0xD1, 0xBE, 0x52, 0x54, 0xD7, 0xA8, 0xF8,
-    0x33, 0xE1, 0xAF, 0xC4, 0xA8, 0x24, 0xD6, 0x6D, 0xED, 0xBC, 0x0A, 0xD3, 0x5A, 0x5C, 0xC4, 0xDF, 0x64, 0x17, 0x97,
-    0x36, 0xF0, 0x65, 0x52, 0x39, 0xD4, 0x23, 0x22, 0x9D, 0xD8, 0x64, 0x83, 0x04, 0x85, 0x38, 0xEF, 0xB8, 0x49, 0x91,
-    0xBB, 0xFF, 0x00, 0x0A, 0xDF, 0xE3, 0x77, 0x84, 0xAC, 0xD3, 0x4A, 0xBC, 0xF0, 0x3A, 0x5A, 0x5B, 0x4F, 0x72, 0x5E,
-    0xD7, 0x53, 0x8E, 0xE2, 0x12, 0x6D, 0xC6, 0x40, 0x20, 0x47, 0xE6, 0x3E, 0x71, 0x26, 0x70, 0x24, 0xE4, 0xB4, 0x63,
-    0x2A, 0xC0, 0xE0, 0x7E, 0x71, 0x98, 0x71, 0x35, 0x5C, 0x4E, 0x69, 0x5A, 0xBD, 0x7C, 0x7A, 0x9C, 0xE3, 0xC9, 0xCF,
-    0x38, 0x46, 0x8D, 0xA4, 0xDC, 0x63, 0x38, 0x4A, 0xE9, 0xED, 0x5A, 0x4E, 0x4A, 0x6D, 0xF3, 0xCA, 0x2D, 0xB9, 0x4D,
-    0xD3, 0x77, 0x8C, 0x7F, 0xA8, 0xF8, 0x53, 0x8D, 0x31, 0x32, 0xC8, 0xA9, 0xDB, 0x15, 0x04, 0x9F, 0xB4, 0xD2, 0xF4,
-    0xD3, 0xF8, 0xA5, 0x77, 0xCD, 0xB3, 0x49, 0x6C, 0xB4, 0xEA, 0x97, 0x46, 0x78, 0x45, 0xC4, 0x16, 0x5A, 0xDD, 0xA5,
-    0xA6, 0x8D, 0xA7, 0xEA, 0x77, 0x4A, 0x54, 0x47, 0x3C, 0x31, 0x41, 0x20, 0x06, 0x68, 0x00, 0x9A, 0x4F, 0x2A, 0x40,
-    0x5D, 0x5F, 0x69, 0xF3, 0x15, 0xCB, 0x32, 0x96, 0xDC, 0xC0, 0x90, 0xA5, 0xF0, 0xBD, 0xAF, 0xC3, 0x96, 0x83, 0x51,
-    0xB8, 0xD6, 0xA7, 0x48, 0xD2, 0xE8, 0x2B, 0x2C, 0x2D, 0x6B, 0xF6, 0xA4, 0x55, 0x8E, 0x45, 0x5F, 0xBA, 0x54, 0x2A,
-    0x36, 0x77, 0x12, 0x72, 0xE5, 0x9C, 0x14, 0x3B, 0x63, 0x6C, 0x62, 0x8A, 0x2B, 0xDE, 0xCD, 0x70, 0x74, 0xA7, 0x92,
-    0xE2, 0x31, 0x35, 0x2D, 0x29, 0x54, 0xBB, 0x97, 0xBB, 0x18, 0xF3, 0x38, 0x52, 0xA5, 0x38, 0xDF, 0x92, 0x31, 0x7E,
-    0xEB, 0x9B, 0x6A, 0xCD, 0x6B, 0xEF, 0x7C, 0x6E, 0x52, 0x7F, 0xE1, 0xF7, 0x0E, 0x54, 0x96, 0x1B, 0x3D, 0x52, 0x86,
-    0xF0, 0x73, 0x4B, 0xEE, 0x92, 0xD6, 0xD6, 0xBF, 0xC2, 0x8D, 0x9F, 0x0B, 0xDE, 0xEB, 0xF3, 0x69, 0x6B, 0x06, 0x9B,
-    0xA9, 0x5B, 0xE9, 0xF6, 0x6F, 0x76, 0x1E, 0x45, 0x8E, 0xE1, 0x54, 0xC8, 0xCD, 0xF3, 0x28, 0x8B, 0x1C, 0x9E, 0x25,
-    0x55, 0xDA, 0xA5, 0x58, 0x80, 0xCA, 0x46, 0x42, 0xA8, 0xF4, 0x1B, 0xEB, 0xBB, 0xBB, 0x8B, 0x69, 0xAF, 0xA5, 0x8E,
-    0xD0, 0x49, 0x25, 0x9C, 0x37, 0x73, 0xE9, 0xD7, 0xF2, 0x14, 0x8E, 0x55, 0x91, 0x64, 0x90, 0x46, 0xA8, 0xF9, 0x50,
-    0x57, 0x73, 0x2E, 0xE0, 0x55, 0xF6, 0x87, 0x45, 0x20, 0x2E, 0xC7, 0x28, 0xAF, 0xCD, 0x78, 0xA6, 0xA5, 0x5C, 0x26,
-    0x33, 0x01, 0x35, 0x2E, 0x67, 0x52, 0x35, 0xA0, 0xEE, 0x93, 0x56, 0x70, 0x4F, 0x48, 0xDB, 0x95, 0x6B, 0x76, 0xF4,
-    0xB4, 0xA5, 0x27, 0x29, 0x29, 0x35, 0x16, 0xBF, 0xA8, 0x7C, 0x3E, 0x92, 0xF6, 0x98, 0xA8, 0x5B, 0x4B, 0xC5, 0xDB,
-    0x5E, 0x8E, 0x4F, 0x7D, 0xFA, 0xDB, 0x57, 0xB2, 0x56, 0xD6, 0xED, 0xF7, 0x3E, 0x10, 0xB2, 0x82, 0xF6, 0xC2, 0xF2,
-    0xF3, 0x4E, 0xD4, 0xAF, 0xEE, 0x5A, 0xCA, 0xCE, 0xD2, 0xF1, 0xA7, 0xB8, 0x9D, 0xD4, 0x42, 0x8C, 0xF2, 0x3C, 0x8E,
-    0x42, 0x02, 0x24, 0xFD, 0xDA, 0x96, 0x60, 0xAD, 0xB0, 0xC9, 0x0E, 0x00, 0xC8, 0xC8, 0xEA, 0xEE, 0x12, 0x7B, 0x09,
-    0x6C, 0x6E, 0xCE, 0xA2, 0x0B, 0xAA, 0x32, 0xD9, 0xB9, 0x49, 0x66, 0xDF, 0x74, 0xCC, 0x31, 0x1B, 0xE4, 0xE6, 0x63,
-    0xF2, 0x91, 0x95, 0x42, 0x15, 0x8B, 0x7D, 0xE0, 0xAC, 0xB4, 0x51, 0x5F, 0x33, 0x3A, 0xD8, 0x8C, 0x46, 0x7A, 0x9C,
-    0xA6, 0xD7, 0x2C, 0x13, 0x49, 0x25, 0x6E, 0x6A, 0x73, 0xE4, 0x84, 0x9E, 0x9A, 0xC9, 0x29, 0x24, 0x9E, 0xE9, 0x42,
-    0x11, 0x8D, 0xA2, 0x9C, 0x5F, 0xF5, 0x37, 0x0D, 0x36, 0xF8, 0x7E, 0x9A, 0xEC, 0xE5, 0xE5, 0x7E, 0x67, 0x77, 0x7B,
-    0x7A, 0xBF, 0x5E, 0xB7, 0x3F, 0xFF, 0xD9,
-};
-
 static void set_loop_time(uint32_t now) { LoopBlockingGuard clock(nullptr, LOG_STR("test"), now); }
 
 class OnlineImageDownloadEnd : public ::testing::Test {
  protected:
-  std::unique_ptr<OnlineImage> make_image_(runtime_image::ImageFormat format, uint32_t buffer_size = 256) {
+  std::unique_ptr<OnlineImage> make_image_(runtime_image::ImageFormat format) {
     auto img = std::make_unique<OnlineImage>("http://example.com/image", 0, 0, format, image::IMAGE_TYPE_RGB,
-                                             image::TRANSPARENCY_OPAQUE, nullptr, buffer_size);
+                                             image::TRANSPARENCY_OPAQUE, nullptr, 256);
     img->set_parent(&this->http_);
     return img;
   }
@@ -458,294 +286,4 @@ TEST_F(OnlineImageDownloadEnd, JpegBufferAllocationFailureReportsError) {
   EXPECT_EQ(this->finished_, 0);
 }
 
-// The JPEG decoded with all of it in memory at once, to compare the downloaded images with
-static std::vector<uint8_t> reference_jpeg_pixels() {
-  runtime_image::RuntimeImage ref(runtime_image::JPEG, image::IMAGE_TYPE_RGB, image::TRANSPARENCY_OPAQUE);
-  std::vector<uint8_t> data(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40));
-  EXPECT_TRUE(ref.begin_decode(data.size()));
-  EXPECT_EQ(ref.feed_data(data.data(), data.size()), static_cast<int>(data.size()));
-  EXPECT_TRUE(ref.end_decode());
-  return std::vector<uint8_t>(ref.get_data_start(), ref.get_data_start() + ref.get_width_stride() * ref.get_height());
-}
-
-static std::vector<uint8_t> image_pixels(OnlineImage &image) {
-  return std::vector<uint8_t>(image.get_data_start(),
-                              image.get_data_start() + image.get_width_stride() * image.get_height());
-}
-
-// A JPEG of known size is decoded as it arrives, so a buffer smaller than the file only has to hold the headers
-TEST_F(OnlineImageDownloadEnd, JpegDecodesThroughASmallBuffer) {
-  auto image = this->make_image_(runtime_image::JPEG, 1024);
-  // Bytes after the image are never needed. Waiting for the whole file before decoding would read them all.
-  std::vector<uint8_t> body(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40));
-  body.resize(body.size() + 8192, 0);
-  auto container = std::make_shared<FakeContainer>(body, body.size(), false);
-  container->add_response_header("etag", "\"abc\"");
-  this->start_(*image, container);
-
-  for (int i = 0; i < 100 && this->finished_ + this->errors_ == 0; i++)
-    image->loop();
-
-  ASSERT_EQ(this->finished_, 1);
-  EXPECT_LT(container->get_bytes_read(), body.size()) << "the image was not decoded as it arrived";
-  EXPECT_EQ(this->errors_, 0);
-  ASSERT_EQ(image->get_width(), 48);
-  ASSERT_EQ(image->get_height(), 40);
-  EXPECT_EQ(image_pixels(*image), reference_jpeg_pixels());
-
-  // A complete image is cached
-  this->http_.next = std::make_shared<FakeContainer>(std::vector<uint8_t>{}, 0, false);
-  image->update();
-  EXPECT_TRUE(this->http_.sent_header("if-none-match"));
-}
-
-// A read() waits for every byte it asks for, so a large buffer is filled a little at a time rather than in one
-// read that blocks the loop until the whole buffer has arrived
-TEST_F(OnlineImageDownloadEnd, JpegReadsInSmallPiecesWithALargeBuffer) {
-  auto image = this->make_image_(runtime_image::JPEG, 65536);
-  auto container =
-      std::make_shared<FakeContainer>(std::vector<uint8_t>(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40)), 0, false);
-  this->start_(*image, container);
-
-  for (int i = 0; i < 100 && this->finished_ + this->errors_ == 0; i++) {
-    container->deliver(sizeof(JPEG_48X40));
-    image->loop();
-  }
-
-  ASSERT_EQ(this->finished_, 1);
-  EXPECT_EQ(this->errors_, 0);
-  EXPECT_LE(container->largest_read, READ_CHUNK);
-  EXPECT_EQ(image_pixels(*image), reference_jpeg_pixels());
-}
-
-// A downloaded PNG is read a little at a time too
-TEST_F(OnlineImageDownloadEnd, PngReadsInSmallPiecesWithALargeBuffer) {
-  auto image = this->make_image_(runtime_image::PNG, 65536);
-  std::vector<uint8_t> png(PNG_1X1.begin(), PNG_1X1.end());
-  auto container = std::make_shared<FakeContainer>(png, png.size(), false);
-  this->start_(*image, container);
-
-  for (int i = 0; i < 100 && this->finished_ + this->errors_ == 0; i++)
-    image->loop();
-
-  ASSERT_EQ(this->finished_, 1);
-  EXPECT_EQ(this->errors_, 0);
-  EXPECT_LE(container->largest_read, READ_CHUNK);
-}
-
-// A slow link: the decoder waits part way through the image for the rest to arrive
-TEST_F(OnlineImageDownloadEnd, JpegWaitsForASlowDownload) {
-  auto image = this->make_image_(runtime_image::JPEG, 1024);
-  auto container =
-      std::make_shared<FakeContainer>(std::vector<uint8_t>(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40)), 1024, false);
-  this->start_(*image, container);
-
-  int passes = 0;
-  for (; passes < 1000 && this->finished_ + this->errors_ == 0; passes++) {
-    image->loop();
-    if (passes == 5) {
-      EXPECT_EQ(this->finished_, 0) << "the image cannot be finished before the rest of it arrives";
-    }
-    container->deliver(40);
-  }
-
-  ASSERT_EQ(this->finished_, 1);
-  EXPECT_EQ(this->errors_, 0);
-  EXPECT_GT(passes, 10);
-  EXPECT_EQ(image_pixels(*image), reference_jpeg_pixels());
-}
-
-// A download that stops part way through times out rather than waiting for ever
-TEST_F(OnlineImageDownloadEnd, JpegDownloadThatStopsTimesOut) {
-  auto image = this->make_image_(runtime_image::JPEG, 1024);
-  auto container =
-      std::make_shared<FakeContainer>(std::vector<uint8_t>(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40)), 1500, false);
-  set_loop_time(1000);
-  this->start_(*image, container);
-
-  for (int i = 0; i < 10; i++)
-    image->loop();
-  EXPECT_EQ(this->errors_, 0) << "still within the timeout";
-
-  set_loop_time(1000 + this->http_.get_timeout());
-  image->loop();
-  EXPECT_EQ(this->errors_, 1);
-  EXPECT_EQ(this->finished_, 0);
-}
-
-// Headers that do not fit the download buffer cannot be read as the image arrives, so the whole image is
-// downloaded first, keeping the bytes already read
-TEST_F(OnlineImageDownloadEnd, JpegHeadersLargerThanTheBufferDownloadTheWholeImage) {
-  auto image = this->make_image_(runtime_image::JPEG, 256);
-  auto container = std::make_shared<FakeContainer>(std::vector<uint8_t>(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40)),
-                                                   sizeof(JPEG_48X40), false);
-  this->start_(*image, container);
-
-  for (int i = 0; i < 100 && this->finished_ + this->errors_ == 0; i++) {
-    image->loop();
-    container->deliver(256);
-  }
-
-  ASSERT_EQ(this->finished_, 1);
-  EXPECT_EQ(this->errors_, 0);
-  EXPECT_EQ(container->get_bytes_read(), sizeof(JPEG_48X40));
-  EXPECT_EQ(image_pixels(*image), reference_jpeg_pixels());
-}
-
-// The pass that opens the image does not decode any of it
-TEST_F(OnlineImageDownloadEnd, JpegOpensInAPassOfItsOwn) {
-  auto image = this->make_image_(runtime_image::JPEG, 1024);
-  auto container = std::make_shared<FakeContainer>(std::vector<uint8_t>(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40)),
-                                                   sizeof(JPEG_48X40), false);
-  this->start_(*image, container);
-
-  image->loop();
-  EXPECT_EQ(this->finished_ + this->errors_, 0);
-  for (int i = 0; i < 100 && this->finished_ + this->errors_ == 0; i++)
-    image->loop();
-  EXPECT_EQ(this->finished_, 1);
-}
-
-// A progressive JPEG is refused while it downloads, without blaming the buffer size
-TEST_F(OnlineImageDownloadEnd, JpegProgressiveReportsError) {
-  auto image = this->make_image_(runtime_image::JPEG, 256);
-  std::vector<uint8_t> body(JPEG_PROGRESSIVE, JPEG_PROGRESSIVE + sizeof(JPEG_PROGRESSIVE));
-  auto container = std::make_shared<FakeContainer>(body, 256, false);
-  this->start_(*image, container);
-
-  for (int i = 0; i < 10 && this->finished_ + this->errors_ == 0; i++)
-    image->loop();
-
-  EXPECT_EQ(this->errors_, 1);
-  EXPECT_EQ(this->finished_, 0);
-}
-
-// on_download_started fires once the server has accepted the request, before the first progress report, and not
-// for a response that has no image to download
-TEST_F(OnlineImageDownloadEnd, StartedFiresOnlyWhenTheDownloadBegins) {
-  auto image = this->make_image_(runtime_image::JPEG, 1024);
-  std::vector<std::string> events;
-  image->add_on_download_started_callback([&events]() { events.emplace_back("started"); });
-  image->add_on_progress_callback([&events](float, size_t) { events.emplace_back("progress"); });
-  this->start_(*image,
-               std::make_shared<FakeContainer>(std::vector<uint8_t>(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40)),
-                                               sizeof(JPEG_48X40), false));
-  ASSERT_GE(events.size(), 2u);
-  EXPECT_EQ(events[0], "started");
-  EXPECT_EQ(events[1], "progress");
-  for (int i = 0; i < 100 && this->finished_ + this->errors_ == 0; i++)
-    image->loop();
-  EXPECT_EQ(std::count(events.begin(), events.end(), std::string("started")), 1);
-
-  for (int status : {304, 404}) {
-    events.clear();
-    auto container = std::make_shared<FakeContainer>(std::vector<uint8_t>{}, 0, false);
-    container->status_code = status;
-    this->http_.next = container;
-    image->update();
-    EXPECT_TRUE(events.empty()) << "HTTP " << status;
-  }
-}
-
-// Progress starts at 0, is reported at most every 100 ms, and always reaches 100%
-TEST_F(OnlineImageDownloadEnd, ProgressIsReportedAsTheImageDownloads) {
-  auto image = this->make_image_(runtime_image::JPEG, 1024);
-  std::vector<float> reports;
-  image->add_on_progress_callback([&reports](float progress, size_t) { reports.push_back(progress); });
-  auto container =
-      std::make_shared<FakeContainer>(std::vector<uint8_t>(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40)), 100, false);
-  uint32_t now = 1000;
-  set_loop_time(now);
-  this->start_(*image, container);
-  ASSERT_EQ(reports.size(), 1u);
-  EXPECT_EQ(reports[0], 0.0f);
-
-  int passes = 0;
-  for (; passes < 1000 && this->finished_ + this->errors_ == 0; passes++) {
-    now += 50;
-    set_loop_time(now);
-    container->deliver(100);
-    image->loop();
-  }
-
-  ASSERT_EQ(this->finished_, 1);
-  EXPECT_EQ(reports.back(), 100.0f);
-  EXPECT_TRUE(std::is_sorted(reports.begin(), reports.end()));
-  // A report every 100 ms of a 50 ms loop, plus the start and the end
-  EXPECT_LE(reports.size(), static_cast<size_t>(passes / 2 + 2));
-  EXPECT_GT(reports.size(), 3u);
-}
-
-// Without a Content-Length the percentage is unknown, but the bytes read are still reported
-TEST_F(OnlineImageDownloadEnd, ProgressOfAChunkedDownloadHasNoPercentage) {
-  auto image = this->make_image_(runtime_image::PNG);
-  std::vector<std::pair<float, size_t>> reports;
-  image->add_on_progress_callback(
-      [&reports](float progress, size_t bytes_read) { reports.emplace_back(progress, bytes_read); });
-  auto container =
-      std::make_shared<FakeContainer>(std::vector<uint8_t>(PNG_1X1.begin(), PNG_1X1.end()), PNG_1X1.size(), true);
-  set_loop_time(1000);
-  this->start_(*image, container);
-  set_loop_time(2000);
-  for (int i = 0; i < 10 && this->finished_ + this->errors_ == 0; i++)
-    image->loop();
-
-  ASSERT_GE(reports.size(), 2u);
-  EXPECT_TRUE(std::isnan(reports.back().first));
-  EXPECT_EQ(reports.back().second, PNG_1X1.size());
-}
-
-// Scan data that cannot be decoded fails the download part way through
-TEST_F(OnlineImageDownloadEnd, JpegDecodeErrorFailsTheDownload) {
-  auto image = this->make_image_(runtime_image::JPEG, 1024);
-  std::vector<uint8_t> body(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40));
-  std::fill(body.begin() + 700, body.end() - 2, 0xFE);
-  auto container = std::make_shared<FakeContainer>(body, 1024, false);
-  this->start_(*image, container);
-
-  for (int i = 0; i < 100 && this->finished_ + this->errors_ == 0; i++) {
-    image->loop();
-    container->deliver(256);
-  }
-
-  EXPECT_EQ(this->errors_, 1);
-  EXPECT_EQ(this->finished_, 0);
-}
-
-// A read that fails part way through the download ends it with an error
-TEST_F(OnlineImageDownloadEnd, JpegReadErrorFailsTheDownload) {
-  auto image = this->make_image_(runtime_image::JPEG, 1024);
-  auto container =
-      std::make_shared<FakeContainer>(std::vector<uint8_t>(JPEG_48X40, JPEG_48X40 + sizeof(JPEG_48X40)), 1024, false);
-  this->start_(*image, container);
-  for (int i = 0; i < 5; i++)
-    image->loop();
-  ASSERT_EQ(this->errors_, 0);
-
-  container->fail_reads();
-  image->loop();
-  EXPECT_EQ(this->errors_, 1);
-  EXPECT_EQ(this->finished_, 0);
-}
-
-TEST(DownloadBufferGrow, KeepsTheUnreadData) {
-  DownloadBuffer buffer(16);
-  std::memcpy(buffer.append(), "abcdef", 6);
-  buffer.write(6);
-  buffer.read(2);
-
-  ASSERT_TRUE(buffer.grow(64));
-  EXPECT_EQ(buffer.size(), 64u);
-  ASSERT_EQ(buffer.unread(), 4u);
-  EXPECT_EQ(std::memcmp(buffer.data(), "cdef", 4), 0);
-
-  // A failed allocation leaves the buffer as it was
-  EXPECT_FALSE(buffer.grow(SIZE_MAX / 2));
-  EXPECT_EQ(buffer.size(), 64u);
-  EXPECT_EQ(std::memcmp(buffer.data(), "cdef", 4), 0);
-  // Growing to a smaller size changes nothing
-  EXPECT_TRUE(buffer.grow(8));
-  EXPECT_EQ(buffer.size(), 64u);
-}
-
 }  // namespace esphome::online_image::testing

tests/components/runtime_image/test_decoder_reuse.cpp

@@ -3,13 +3,10 @@
 
 #include <algorithm>
 #include <array>
-#include <cstdlib>
 #include <cstring>
-#include <utility>
 #include <vector>
 
 #include "esphome/components/runtime_image/image_decoder.h"
-#include "esphome/components/runtime_image/jpeg_decoder.h"
 #include "esphome/components/runtime_image/runtime_image.h"
 
 namespace esphome::runtime_image::testing {
@@ -365,23 +362,6 @@ static std::vector<uint8_t> pixel_bytes(TestableRuntimeImage &img) {
   return std::vector<uint8_t>(start, start + img.get_width_stride() * img.get_height());
 }
 
-// Shrinking a PNG keeps the first source pixel of each span, as the JPEG path does: the even rows and columns
-TEST(RuntimeImageDecoder, ShrunkPngKeepsTheFirstPixelOfEachSpan) {
-  TestableRuntimeImage img(PNG, image::IMAGE_TYPE_RGB, false, 2, 2);
-  ASSERT_TRUE(decode_all(img, PNG_RGB, sizeof(PNG_RGB), PNG));
-  ASSERT_EQ(img.get_width(), 2);
-  ASSERT_EQ(img.get_height(), 2);
-  for (int row = 0; row < 2; row++) {
-    for (int col = 0; col < 2; col++) {
-      const Color pixel = img.get_pixel(col, row);
-      const uint8_t *expected = PNG_RGB_EXPECTED[row * 2][col * 2];
-      EXPECT_EQ(pixel.r, expected[0]) << col << "," << row;
-      EXPECT_EQ(pixel.g, expected[1]) << col << "," << row;
-      EXPECT_EQ(pixel.b, expected[2]) << col << "," << row;
-    }
-  }
-}
-
 TEST(RuntimeImageDecoder, JpegDecoderStaysWarmAcrossDecodes) {
   TestableRuntimeImage img(JPEG);
 
@@ -397,12 +377,6 @@ TEST(RuntimeImageDecoder, JpegDecoderStaysWarmAcrossDecodes) {
   EXPECT_EQ(pixel_bytes(img), first_pixels) << "reused decoder must reproduce identical pixels";
 }
 
-// The RGB565 word at (x, y) of a decoded image, in its own byte order
-static uint16_t rgb565_word(const TestableRuntimeImage &img, int x, int y) {
-  const uint8_t *p = img.get_data_start() + (y * img.get_width() + x) * 2;
-  return p[0] | (p[1] << 8);
-}
-
 // JPEGDEC's colour conversion differs by one step between its SIMD and scalar builds, so each
 // channel is checked within one step; the byte order is still exact
 static void expect_rgb565_near(uint16_t actual, uint16_t expected, const char *order, size_t pixel) {
@@ -412,12 +386,16 @@ static void expect_rgb565_near(uint16_t actual, uint16_t expected, const char *o
 }
 
 TEST(RuntimeImageDecoder, JpegDecodesDirectlyToScaledRgb565WithBothByteOrders) {
-  // Halving the 8x8 gradient makes JPEGDEC decode it at half size, which averages each 2x2 block of pixels
-  TestableRuntimeImage full(JPEG, image::IMAGE_TYPE_RGB565, false, 8, 8);
+  // Halving the 8x8 gradient keeps its even rows and columns. The words are JPEGDEC's output for those
+  // pixels, red in the top five bits: a decode of the fixture with PIL gives (51, 7, 30) for source
+  // pixel (2, 0), which is word 0x3803 here.
+  static const uint16_t EXPECTED[4 * 4] = {
+      0x0000, 0x3803, 0x8008, 0xB80B, 0x0203, 0x3A07, 0x81EC, 0xBA0F,
+      0x03E7, 0x3C0B, 0x83F0, 0xBC13, 0x05EC, 0x3E0F, 0x85F4, 0xBE17,
+  };
   TestableRuntimeImage rgb565_little_endian(JPEG, image::IMAGE_TYPE_RGB565, false, 4, 4);
   TestableRuntimeImage rgb565_big_endian(JPEG, image::IMAGE_TYPE_RGB565, true, 4, 4);
 
-  ASSERT_TRUE(decode_all(full, JPEG_GRADIENT, sizeof(JPEG_GRADIENT)));
   ASSERT_TRUE(decode_all(rgb565_little_endian, JPEG_GRADIENT, sizeof(JPEG_GRADIENT)));
   ASSERT_TRUE(decode_all(rgb565_big_endian, JPEG_GRADIENT, sizeof(JPEG_GRADIENT)));
   ASSERT_EQ(rgb565_little_endian.get_width(), 4);
@@ -427,45 +405,9 @@ TEST(RuntimeImageDecoder, JpegDecodesDirectlyToScaledRgb565WithBothByteOrders) {
 
   const uint8_t *little = rgb565_little_endian.get_data_start();
   const uint8_t *big = rgb565_big_endian.get_data_start();
-  for (int row = 0; row < 4; row++) {
-    for (int col = 0; col < 4; col++) {
-      int red = 0;
-      int green = 0;
-      int blue = 0;
-      for (int dy = 0; dy < 2; dy++) {
-        for (int dx = 0; dx < 2; dx++) {
-          const uint16_t word = rgb565_word(full, col * 2 + dx, row * 2 + dy);
-          red += word >> 11;
-          green += (word >> 5) & 0x3F;
-          blue += word & 0x1F;
-        }
-      }
-      const uint16_t expected = ((red + 2) / 4) << 11 | ((green + 2) / 4) << 5 | (blue + 2) / 4;
-      const size_t i = row * 4 + col;
-      expect_rgb565_near(little[i * 2] | (little[i * 2 + 1] << 8), expected, "little-endian", i);
-      expect_rgb565_near((big[i * 2] << 8) | big[i * 2 + 1], expected, "big-endian", i);
-    }
-  }
-}
-
-// The generic draw() path gets JPEGDEC's half size RGB8888 output. It must match the half size RGB565 decode checked
-// above. A full size RGB8888 decode is no reference: JPEGDEC's x86 SSE code writes it with red and blue swapped.
-TEST(RuntimeImageDecoder, JpegDecodesToScaledRgb) {
-  TestableRuntimeImage rgb565(JPEG, image::IMAGE_TYPE_RGB565, false, 4, 4);
-  TestableRuntimeImage rgb(JPEG, image::IMAGE_TYPE_RGB, false, 4, 4);
-  ASSERT_TRUE(decode_all(rgb565, JPEG_GRADIENT, sizeof(JPEG_GRADIENT)));
-  ASSERT_TRUE(decode_all(rgb, JPEG_GRADIENT, sizeof(JPEG_GRADIENT)));
-  ASSERT_EQ(static_cast<JpegDecoder *>(rgb.decoder())->get_scale(), 2);
-
-  for (int row = 0; row < 4; row++) {
-    for (int col = 0; col < 4; col++) {
-      const uint16_t word = rgb565_word(rgb565, col, row);
-      const Color pixel = rgb.get_pixel(col, row);
-      // RGB565 keeps the top 5 or 6 bits of each channel
-      EXPECT_LE(std::abs(pixel.r - ((word >> 11) << 3)), 8) << "pixel " << col << "," << row << " red";
-      EXPECT_LE(std::abs(pixel.g - (((word >> 5) & 0x3F) << 2)), 4) << "pixel " << col << "," << row << " green";
-      EXPECT_LE(std::abs(pixel.b - ((word & 0x1F) << 3)), 8) << "pixel " << col << "," << row << " blue";
-    }
+  for (size_t i = 0; i < 4 * 4; i++) {
+    expect_rgb565_near(little[i * 2] | (little[i * 2 + 1] << 8), EXPECTED[i], "little-endian", i);
+    expect_rgb565_near((big[i * 2] << 8) | big[i * 2 + 1], EXPECTED[i], "big-endian", i);
   }
 }
 
@@ -473,43 +415,10 @@ TEST(RuntimeImageDecoder, JpegDecodesToScaledRgb) {
 // prepare_rgb565_scaling_ rounds: source n covers targets ceil(n * target / source) up to the next start
 static uint16_t scaled_boundary(int n, int source, int target) { return (n * target + source - 1) / source; }
 
-// Every target pixel of a PNG drawn at a size that is not a whole multiple: enlarged with spans 1 or 2 pixels wide, and
-// shrunk with spans of 0, 1 or 2 pixels
-TEST(RuntimeImageDecoder, ScaledPngPaintsEachSourcePixelOverItsSpan) {
-  // A fixed size keeps the aspect ratio, so both are square
-  for (auto [width, height] : {std::pair{6, 6}, std::pair{3, 3}}) {
-    TestableRuntimeImage img(PNG, image::IMAGE_TYPE_RGB, false, width, height);
-    ASSERT_TRUE(decode_all(img, PNG_RGB, sizeof(PNG_RGB), PNG));
-    ASSERT_EQ(img.get_width(), width);
-    ASSERT_EQ(img.get_height(), height);
-    int painted = 0;
-    for (int m = 0; m < 4; m++) {
-      for (int n = 0; n < 4; n++) {
-        const uint8_t *expected = PNG_RGB_EXPECTED[m][n];
-        for (int y = scaled_boundary(m, 4, height); y < scaled_boundary(m + 1, 4, height); y++) {
-          for (int x = scaled_boundary(n, 4, width); x < scaled_boundary(n + 1, 4, width); x++) {
-            const Color pixel = img.get_pixel(x, y);
-            EXPECT_EQ(pixel.r, expected[0]) << width << "x" << height << " at " << x << "," << y;
-            EXPECT_EQ(pixel.g, expected[1]) << width << "x" << height << " at " << x << "," << y;
-            EXPECT_EQ(pixel.b, expected[2]) << width << "x" << height << " at " << x << "," << y;
-            painted++;
-          }
-        }
-      }
-    }
-    EXPECT_EQ(painted, width * height);
-  }
-}
-
-// JPEGDEC decodes at the smallest of full, 1/2, 1/4 and 1/8 size that still covers the buffer
-TEST(RuntimeImageDecoder, JpegDecodesAtTheSmallestSizeThatCoversTheBuffer) {
-  for (auto [size, scale] :
-       {std::pair{8, 1}, std::pair{5, 1}, std::pair{4, 2}, std::pair{3, 2}, std::pair{2, 4}, std::pair{1, 8}}) {
-    TestableRuntimeImage img(JPEG, image::IMAGE_TYPE_RGB565, false, size, size);
-    ASSERT_TRUE(decode_all(img, JPEG_GRADIENT, sizeof(JPEG_GRADIENT)));
-    EXPECT_EQ(img.get_width(), size);
-    EXPECT_EQ(static_cast<JpegDecoder *>(img.decoder())->get_scale(), scale) << size << "x" << size;
-  }
+// The RGB565 word at (x, y) of a decoded image, in its own byte order
+static uint16_t rgb565_word(const TestableRuntimeImage &img, int x, int y) {
+  const uint8_t *p = img.get_data_start() + (y * img.get_width() + x) * 2;
+  return p[0] | (p[1] << 8);
 }
 
 // Every scaled direct decode must equal the unscaled one sampled through the boundary table,
@@ -518,8 +427,7 @@ TEST(RuntimeImageDecoder, JpegDirectRgb565ScalesBySourcePixelSpans) {
   TestableRuntimeImage full(JPEG, image::IMAGE_TYPE_RGB565, false, 8, 8);
   ASSERT_TRUE(decode_all(full, JPEG_GRADIENT, sizeof(JPEG_GRADIENT)));
 
-  // Sizes JPEGDEC cannot reach by scaling, so the full size decode is spread over the buffer
-  for (int size : {5, 6, 16}) {
+  for (int size : {4, 6, 16}) {
     TestableRuntimeImage scaled(JPEG, image::IMAGE_TYPE_RGB565, false, size, size);
     ASSERT_TRUE(decode_all(scaled, JPEG_GRADIENT, sizeof(JPEG_GRADIENT)));
     ASSERT_EQ(scaled.get_width(), size);

tests/components/runtime_image/test_grayscale_pixel.cpp

@@ -1,35 +0,0 @@
-#include <gtest/gtest.h>
-
-#include <cstdint>
-#include <cstdlib>
-
-#include "esphome/components/runtime_image/runtime_image.h"
-
-namespace esphome::runtime_image::testing {
-
-class GrayscaleImage : public RuntimeImage {
- public:
-  GrayscaleImage() : RuntimeImage(BMP, image::IMAGE_TYPE_GRAYSCALE, image::TRANSPARENCY_OPAQUE) {}
-  uint8_t pixel(int x) const { return this->buffer_[x]; }
-};
-
-// The integer weights stay within one step of the floating point formula they replace, and keep white at 255
-TEST(GrayscalePixel, MatchesTheFloatingPointWeights) {
-  GrayscaleImage img;
-  ASSERT_GT(img.resize(256, 1), 0);
-  for (int r = 0; r < 256; r += 5) {
-    for (int g = 0; g < 256; g += 5) {
-      for (int b = 0; b < 256; b += 5) {
-        img.draw_pixel(0, 0, Color(r, g, b));
-        const int expected = static_cast<int>(0.2125 * r + 0.7154 * g + 0.0721 * b);
-        ASSERT_LE(std::abs(img.pixel(0) - expected), 1) << "r=" << r << " g=" << g << " b=" << b;
-      }
-    }
-  }
-  img.draw_pixel(1, 0, Color(255, 255, 255));
-  EXPECT_EQ(img.pixel(1), 255);
-  img.draw_pixel(2, 0, Color(0, 0, 0));
-  EXPECT_EQ(img.pixel(2), 0);
-}
-
-}  // namespace esphome::runtime_image::testing

tests/components/runtime_image/test_jpeg_source.cpp

@@ -5,7 +5,6 @@
 #include <vector>
 
 #include "esphome/components/runtime_image/image_decoder.h"
-#include "esphome/components/runtime_image/jpeg_decoder.h"
 #include "esphome/components/runtime_image/runtime_image.h"
 
 #ifdef USE_RUNTIME_IMAGE_JPEG
@@ -179,7 +178,6 @@ class SourceImage : public RuntimeImage {
   explicit SourceImage(image::ImageType type = image::IMAGE_TYPE_RGB, int width = 0, int height = 0)
       : RuntimeImage(JPEG, type, image::TRANSPARENCY_OPAQUE, nullptr, false, width, height) {}
   size_t decoded_bytes() const { return this->decoded_bytes_; }
-  int scale() const { return static_cast<const JpegDecoder *>(this->decoder_.get())->get_scale(); }
 };
 
 /// The file in memory; fails reads once fail_at bytes have been read
@@ -293,23 +291,6 @@ TEST(JpegSourceDecode, SteppedDirectRgb565MatchesFedDecode) {
   EXPECT_EQ(pixels(stepped), pixels(fed));
 }
 
-// decodeBegin() is given the same reduced scale as decode(), and decodeRows() keeps it for every step
-TEST(JpegSourceDecode, SteppedScaledDecodeMatchesFedDecode) {
-  for (auto type : {image::IMAGE_TYPE_RGB565, image::IMAGE_TYPE_RGB}) {
-    // 24x20 is half of 48x40, so JPEGDEC decodes at 1/2 size
-    SourceImage fed(type, 24, 20);
-    ASSERT_TRUE(decode_fed(fed));
-    ASSERT_EQ(fed.scale(), 2);
-
-    SourceImage stepped(type, 24, 20);
-    MemorySource source(JPEG_ROWS, sizeof(JPEG_ROWS));
-    int steps;
-    ASSERT_EQ(decode_stepped(stepped, source, steps), DECODE_STEP_DONE);
-    ASSERT_EQ(stepped.scale(), 2);
-    EXPECT_EQ(pixels(stepped), pixels(fed)) << "image type " << type;
-  }
-}
-
 TEST(JpegSourceDecode, WaitsForDataThatHasNotArrived) {
   SourceImage fed;
   ASSERT_TRUE(decode_fed(fed));

tests/integration/test_addressable_light_transition.py

@@ -23,7 +23,7 @@ import asyncio
 from aioesphomeapi import LightInfo, SensorInfo, SensorState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -52,11 +52,9 @@ async def test_addressable_light_transition(
                 return
             samples.append((loop.time(), state.state))
 
-        # InitialStateHelper swallows the first state ESPHome sends per entity
+        # Swallow the first state ESPHome sends per entity
         # on subscribe, so on_state only sees real post-subscribe updates.
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Start transition: off -> full white over 1 second. This is the
         # scenario from the bug report, compressed in time.

tests/integration/test_alarm_control_panel_state_transitions.py

@@ -15,7 +15,11 @@ from aioesphomeapi import (
 )
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity, require_only_entity
+from .state_utils import (
+    require_entity,
+    require_only_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -99,15 +103,11 @@ async def test_alarm_control_panel_state_transitions(
                 states_received.append(state.state)
                 state_event.set()
 
-        # Use InitialStateHelper to handle initial state broadcast
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for initial states from all entities
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Verify alarm panel started in DISARMED state
-        initial_alarm_state = initial_state_helper.initial_states.get(alarm_info.key)
+        initial_alarm_state = sub.initial_states.get(alarm_info.key)
         assert initial_alarm_state is not None, "No initial alarm state received"
         assert isinstance(initial_alarm_state, AlarmControlPanelEntityState)
         assert initial_alarm_state.state == AlarmControlPanelState.DISARMED
@@ -193,7 +193,7 @@ async def test_alarm_control_panel_state_transitions(
         )
         await wait_for_state(AlarmControlPanelState.DISARMED)
 
-        # Verify basic state sequence (initial DISARMED is handled by InitialStateHelper)
+        # Verify basic state sequence (initial DISARMED is swallowed on subscribe)
         expected_states = [
             AlarmControlPanelState.ARMING,  # Arm away
             AlarmControlPanelState.ARMED_AWAY,

tests/integration/test_api_cover_toggle.py

@@ -18,7 +18,7 @@ from .raw_api_client import (
     decode_fields,
     tag,
 )
-from .state_utils import InitialStateHelper, StateWaiter, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 COVER_COMMAND = MESSAGE_TYPE_OF[api_pb2.CoverCommandRequest]
@@ -48,12 +48,10 @@ async def test_api_cover_toggle(
         plain_cover = require_entity(entities, "plain_cover", CoverInfo, exact=True)
 
         positions: dict[int, list[float]] = {toggle_cover.key: [], plain_cover.key: []}
-        waiter = StateWaiter()
 
         def on_state(state: EntityState) -> None:
             if isinstance(state, CoverState) and state.key in positions:
                 positions[state.key].append(state.position)
-            waiter.on_state(state)
 
         def position_is(key: int, value: float) -> Callable[[EntityState], bool]:
             return lambda s: (
@@ -62,9 +60,8 @@ async def test_api_cover_toggle(
                 and s.position == pytest.approx(value, abs=0.01)
             )
 
-        initial = InitialStateHelper(entities)
-        client.subscribe_states(initial.on_state_wrapper(on_state))
-        await initial.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities, on_state=on_state)
+        waiter = sub.waiter
 
         # Capability discovery: only the cover with a toggle_action advertises toggle
         await raw.connect()

tests/integration/test_api_decode_wire_types.py

@@ -25,7 +25,7 @@ from .raw_api_client import (
     encode_varint,
     tag,
 )
-from .state_utils import InitialStateHelper, StateWaiter, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 SWITCH_COMMAND = MESSAGE_TYPE_OF[api_pb2.SwitchCommandRequest]
@@ -52,12 +52,10 @@ async def test_api_decode_wire_types(
         on, off = tag(2, WIRE_VARINT) + b"\x01", tag(2, WIRE_VARINT) + b"\x00"
 
         switch_states: list[bool] = []
-        waiter = StateWaiter()
 
         def on_state(state: EntityState) -> None:
             if isinstance(state, SwitchState) and state.key == switch.key:
                 switch_states.append(state.state)
-            waiter.on_state(state)
 
         def switch_is(value: bool) -> Callable[[EntityState], bool]:
             return lambda s: (
@@ -69,9 +67,8 @@ async def test_api_decode_wire_types(
                 isinstance(s, NumberState) and s.key == number.key and s.state == value
             )
 
-        initial = InitialStateHelper(entities)
-        client.subscribe_states(initial.on_state_wrapper(on_state))
-        await initial.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities, on_state=on_state)
+        waiter = sub.waiter
         await raw.connect()
 
         # A well formed command: fixed32 key, varint state

tests/integration/test_api_encode_boundaries.py

@@ -14,7 +14,7 @@ from aioesphomeapi import (
 )
 import pytest
 
-from .state_utils import InitialStateHelper, StateWaiter, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 LONG_OPTION = (
@@ -44,16 +44,14 @@ async def test_api_encode_boundaries(
         number = require_entity(entities, "negative_number")
         button = require_entity(entities, "publish_values")
 
-        initial = InitialStateHelper(entities)
-        waiter = StateWaiter()
-        client.subscribe_states(initial.on_state_wrapper(waiter.on_state))
-        await initial.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
+        waiter = sub.waiter
 
         # A float of exactly zero is skipped on the wire and must still read as 0.0, not missing
-        first = initial.initial_states[sensor.key]
+        first = sub.initial_states[sensor.key]
         assert isinstance(first, SensorState)
         assert first.state == 0.0 and not first.missing_state
-        first_number = initial.initial_states[number.key]
+        first_number = sub.initial_states[number.key]
         assert isinstance(first_number, NumberState)
         assert first_number.state == -123.5
 

tests/integration/test_binary_sensor_autorepeat_filter.py

@@ -18,10 +18,10 @@ import asyncio
 import pytest
 
 from .state_utils import (
-    InitialStateHelper,
     SensorStateCollector,
     assert_device_name,
     require_entity,
+    subscribe_after_initial_states,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
@@ -53,13 +53,10 @@ async def test_binary_sensor_autorepeat_filter(
             entities, "release", description="Release button"
         )
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
         )
 
-        await initial_state_helper.wait_for_initial_states()
-
         autorepeat_states = collector.binary_states["autorepeat_sensor"]
 
         # Press: source becomes true, autorepeat passes the initial true through

tests/integration/test_binary_sensor_invalidate_state.py

@@ -9,14 +9,15 @@ Regression test for: https://github.com/esphome/esphome/issues/12252
 
 from __future__ import annotations
 
-import asyncio
-
-from aioesphomeapi import BinarySensorInfo, BinarySensorState, EntityState
+from aioesphomeapi import BinarySensorInfo, BinarySensorState
 import pytest
 
-from .state_utils import InitialStateHelper, assert_device_name, require_entity
+from .state_utils import (
+    assert_device_name,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
-from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -33,19 +34,6 @@ async def test_binary_sensor_invalidate_state(
     3. Publishing false sets missing_state=False and state=False
     4. Invalidating state sets missing_state=True (unknown state)
     """
-    loop = asyncio.get_running_loop()
-
-    # Track state changes
-    states_received: list[BinarySensorState] = []
-    state_future: asyncio.Future[BinarySensorState] = loop.create_future()
-
-    def on_state(state: EntityState) -> None:
-        """Track binary sensor state changes."""
-        if isinstance(state, BinarySensorState):
-            states_received.append(state)
-            if not state_future.done():
-                state_future.set_result(state)
-
     async with (
         run_compiled(yaml_config),
         api_client_connected() as client,
@@ -68,15 +56,11 @@ async def test_binary_sensor_invalidate_state(
             entities, "invalidate", description="Invalidate button"
         )
 
-        # Set up initial state helper to handle the initial state broadcast
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for initial states
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         # Check initial state - should be missing (unknown)
-        initial_state = initial_state_helper.initial_states.get(binary_sensor.key)
+        initial_state = sub.initial_states.get(binary_sensor.key)
         assert initial_state is not None, "No initial state received for binary sensor"
         assert isinstance(initial_state, BinarySensorState)
         assert initial_state.missing_state is True, (
@@ -84,11 +68,9 @@ async def test_binary_sensor_invalidate_state(
         )
 
         # Test 1: Set state to true
-        states_received.clear()
-        state_future = loop.create_future()
+        next_state = sub.waiter.next_state(binary_sensor.key, BinarySensorState)
         client.button_command(set_true_button.key)
-
-        state = await wait_or_fail(state_future, 5.0, "Timeout waiting for state=true")
+        state = await next_state
 
         assert state.missing_state is False, (
             f"After setting true, missing_state should be False, got {state}"
@@ -96,11 +78,9 @@ async def test_binary_sensor_invalidate_state(
         assert state.state is True, f"Expected state=True, got {state}"
 
         # Test 2: Set state to false
-        states_received.clear()
-        state_future = loop.create_future()
+        next_state = sub.waiter.next_state(binary_sensor.key, BinarySensorState)
         client.button_command(set_false_button.key)
-
-        state = await wait_or_fail(state_future, 5.0, "Timeout waiting for state=false")
+        state = await next_state
 
         assert state.missing_state is False, (
             f"After setting false, missing_state should be False, got {state}"
@@ -109,17 +89,10 @@ async def test_binary_sensor_invalidate_state(
 
         # Test 3: Invalidate state (set to unknown)
         # This is the critical test for the bug fix
-        states_received.clear()
-        state_future = loop.create_future()
+        # A timeout here means invalidate_state() is not reported to the API (issue #12252)
+        next_state = sub.waiter.next_state(binary_sensor.key, BinarySensorState)
         client.button_command(invalidate_button.key)
-
-        state = await wait_or_fail(
-            state_future,
-            5.0,
-            "Timeout waiting for invalidated state - "
-            "binary_sensor.invalidate_state() may not be reporting to the API. "
-            "See issue #12252.",
-        )
+        state = await next_state
 
         assert state.missing_state is True, (
             f"After invalidate_state(), missing_state should be True (unknown), "

tests/integration/test_climate_control_action.py

@@ -7,18 +7,12 @@ multiple field combinations and the lambda path.
 
 from __future__ import annotations
 
-import asyncio
+from functools import partial
 
-from aioesphomeapi import (
-    ButtonInfo,
-    ClimateInfo,
-    ClimateMode,
-    ClimateState,
-    EntityState,
-)
+from aioesphomeapi import ClimateInfo, ClimateMode, ClimateState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import press_and_wait, require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -29,56 +23,34 @@ async def test_climate_control_action(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test climate ControlAction with constants and lambdas."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
-        climate_state_future: asyncio.Future[ClimateState] | None = None
-
-        def on_state(state: EntityState) -> None:
-            if (
-                isinstance(state, ClimateState)
-                and climate_state_future is not None
-                and not climate_state_future.done()
-            ):
-                climate_state_future.set_result(state)
-
-        async def wait_for_climate_state(timeout: float = 5.0) -> ClimateState:
-            nonlocal climate_state_future
-            climate_state_future = loop.create_future()
-            try:
-                return await asyncio.wait_for(climate_state_future, timeout)
-            finally:
-                climate_state_future = None
-
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         require_entity(entities, "test_climate", ClimateInfo)
 
-        async def press_and_wait(name: str) -> ClimateState:
-            btn = require_entity(entities, name.lower().replace(" ", "_"), ButtonInfo)
-            client.button_command(btn.key)
-            return await wait_for_climate_state()
+        press = partial(
+            press_and_wait, client, entities, sub.waiter, state_type=ClimateState
+        )
 
         # mode only — set HEAT
-        state = await press_and_wait("Set Mode Heat")
+        state = await press("set_mode_heat")
         assert state.mode == ClimateMode.HEAT
 
         # mode + target_temperature_low + target_temperature_high
-        state = await press_and_wait("Set Mode Temps")
+        state = await press("set_mode_temps")
         assert state.mode == ClimateMode.HEAT_COOL
         assert state.target_temperature_low == pytest.approx(19.0, abs=0.5)
         assert state.target_temperature_high == pytest.approx(23.0, abs=0.5)
 
         # target_temperature_low only
-        state = await press_and_wait("Set Low Only")
+        state = await press("set_low_only")
         assert state.target_temperature_low == pytest.approx(17.5, abs=0.5)
 
         # lambda path: target_temperature_high computed at runtime
-        state = await press_and_wait("Lambda High")
+        state = await press("lambda_high")
         assert state.target_temperature_high == pytest.approx(21.5, abs=0.5)
 
         # mode only — turn off via mode
-        state = await press_and_wait("Set Off")
+        state = await press("set_off")
         assert state.mode == ClimateMode.OFF

tests/integration/test_climate_custom_modes.py

@@ -2,15 +2,12 @@
 
 from __future__ import annotations
 
-import asyncio
-
 import aioesphomeapi
-from aioesphomeapi import ClimateInfo, ClimatePreset, EntityState
+from aioesphomeapi import ClimateInfo, ClimatePreset
 import pytest
 
-from .state_utils import InitialStateHelper, require_only_entity
+from .state_utils import require_only_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
-from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -20,42 +17,13 @@ async def test_climate_custom_fan_modes_and_presets(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test that custom presets are properly exposed and can be changed."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
-        states: dict[int, EntityState] = {}
-        super_saver_future: asyncio.Future[EntityState] = loop.create_future()
-        vacation_future: asyncio.Future[EntityState] = loop.create_future()
-
-        def on_state(state: EntityState) -> None:
-            states[state.key] = state
-            if isinstance(state, aioesphomeapi.ClimateState):
-                # Wait for Super Saver preset
-                if (
-                    state.custom_preset == "Super Saver"
-                    and state.target_temperature_low == 20.0
-                    and state.target_temperature_high == 24.0
-                    and not super_saver_future.done()
-                ):
-                    super_saver_future.set_result(state)
-                # Wait for Vacation Mode preset
-                elif (
-                    state.custom_preset == "Vacation Mode"
-                    and state.target_temperature_low == 15.0
-                    and state.target_temperature_high == 18.0
-                    and not vacation_future.done()
-                ):
-                    vacation_future.set_result(state)
-
         # Get entities and set up state synchronization
         entities, services = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         test_climate = require_only_entity(entities, ClimateInfo)
 
-        # Subscribe with the wrapper that filters initial states
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for all initial states to be broadcast
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         # Verify enum presets are exposed (from preset: config map)
         assert ClimatePreset.AWAY in test_climate.supported_presets, (
@@ -76,7 +44,7 @@ async def test_climate_custom_fan_modes_and_presets(
         )
 
         # Get initial state and verify default preset
-        initial_state = initial_state_helper.initial_states.get(test_climate.key)
+        initial_state = sub.initial_states.get(test_climate.key)
         assert initial_state is not None, "Climate initial state not found"
         assert isinstance(initial_state, aioesphomeapi.ClimateState)
         assert initial_state.custom_preset == "Eco Plus", (
@@ -90,13 +58,17 @@ async def test_climate_custom_fan_modes_and_presets(
         )
 
         # Test changing to "Super Saver" custom preset
-        client.climate_command(test_climate.key, custom_preset="Super Saver")
-
-        super_saver_state = await wait_or_fail(
-            super_saver_future,
-            5.0,
-            "Super Saver preset change not received within 5 seconds",
+        super_saver_seen = sub.waiter.expect(
+            lambda s: (
+                isinstance(s, aioesphomeapi.ClimateState)
+                and s.custom_preset == "Super Saver"
+                and s.target_temperature_low == 20.0
+                and s.target_temperature_high == 24.0
+            ),
+            label="Super Saver preset change",
         )
+        client.climate_command(test_climate.key, custom_preset="Super Saver")
+        super_saver_state = await super_saver_seen
 
         assert isinstance(super_saver_state, aioesphomeapi.ClimateState)
         assert super_saver_state.custom_preset == "Super Saver"
@@ -104,13 +76,17 @@ async def test_climate_custom_fan_modes_and_presets(
         assert super_saver_state.target_temperature_high == 24.0
 
         # Test changing to "Vacation Mode" custom preset
-        client.climate_command(test_climate.key, custom_preset="Vacation Mode")
-
-        vacation_state = await wait_or_fail(
-            vacation_future,
-            5.0,
-            "Vacation Mode preset change not received within 5 seconds",
+        vacation_seen = sub.waiter.expect(
+            lambda s: (
+                isinstance(s, aioesphomeapi.ClimateState)
+                and s.custom_preset == "Vacation Mode"
+                and s.target_temperature_low == 15.0
+                and s.target_temperature_high == 18.0
+            ),
+            label="Vacation Mode preset change",
         )
+        client.climate_command(test_climate.key, custom_preset="Vacation Mode")
+        vacation_state = await vacation_seen
 
         assert isinstance(vacation_state, aioesphomeapi.ClimateState)
         assert vacation_state.custom_preset == "Vacation Mode"

tests/integration/test_cover_control_action.py

@@ -7,12 +7,12 @@ implementation. Exercises multiple field combinations and the lambda path.
 
 from __future__ import annotations
 
-import asyncio
+from functools import partial
 
-from aioesphomeapi import ButtonInfo, CoverInfo, CoverState, EntityState
+from aioesphomeapi import CoverInfo, CoverState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import press_and_wait, require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -23,80 +23,58 @@ async def test_cover_control_action(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test cover.control and cover.template.publish with constants and lambdas."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
-        cover_state_future: asyncio.Future[CoverState] | None = None
-
-        def on_state(state: EntityState) -> None:
-            if (
-                isinstance(state, CoverState)
-                and cover_state_future is not None
-                and not cover_state_future.done()
-            ):
-                cover_state_future.set_result(state)
-
-        async def wait_for_cover_state(timeout: float = 5.0) -> CoverState:
-            nonlocal cover_state_future
-            cover_state_future = loop.create_future()
-            try:
-                return await asyncio.wait_for(cover_state_future, timeout)
-            finally:
-                cover_state_future = None
-
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         require_entity(entities, "test_cover", CoverInfo)
 
-        async def press_and_wait(name: str) -> CoverState:
-            btn = require_entity(entities, name.lower().replace(" ", "_"), ButtonInfo)
-            client.button_command(btn.key)
-            return await wait_for_cover_state()
+        press = partial(
+            press_and_wait, client, entities, sub.waiter, state_type=CoverState
+        )
 
         # cover.control: position only
-        state = await press_and_wait("Set Position")
+        state = await press("set_position")
         assert state.position == pytest.approx(0.5, abs=0.01)
 
         # cover.control: tilt only
-        state = await press_and_wait("Set Tilt")
+        state = await press("set_tilt")
         assert state.tilt == pytest.approx(0.75, abs=0.01)
 
         # cover.control: position + tilt
-        state = await press_and_wait("Set Pos Tilt")
+        state = await press("set_pos_tilt")
         assert state.position == pytest.approx(0.25, abs=0.01)
         assert state.tilt == pytest.approx(0.30, abs=0.01)
 
         # cover.control: state alias for position 1.0
-        state = await press_and_wait("Open State")
+        state = await press("open_state")
         assert state.position == pytest.approx(1.0, abs=0.01)
 
         # cover.control: lambda position (test_position global = 0.42)
-        state = await press_and_wait("Lambda Position")
+        state = await press("lambda_position")
         assert state.position == pytest.approx(0.42, abs=0.01)
 
         # cover.template.publish: position only
-        state = await press_and_wait("Publish Pos")
+        state = await press("publish_pos")
         assert state.position == pytest.approx(0.6, abs=0.01)
 
         # cover.template.publish: current_operation only
-        state = await press_and_wait("Publish Op")
+        state = await press("publish_op")
         # CoverOperation.OPENING == 1
         assert state.current_operation == 1
 
         # cover.control: stop only — template cover's stop_action publishes
         # current_operation: IDLE.
-        state = await press_and_wait("Stop Cover")
+        state = await press("stop_cover")
         # CoverOperation.IDLE == 0
         assert state.current_operation == 0
 
         # stop_tilt_action publishes tilt 0.4
-        state = await press_and_wait("Stop Tilt")
+        state = await press("stop_tilt")
         assert state.tilt == pytest.approx(0.4, abs=0.01)
 
-        state = await press_and_wait("Set Tilt")
+        state = await press("set_tilt")
         assert state.tilt == pytest.approx(0.75, abs=0.01)
 
-        state = await press_and_wait("Control Stop Tilt")
+        state = await press("control_stop_tilt")
         assert state.tilt == pytest.approx(0.4, abs=0.01)

tests/integration/test_cover_stop_tilt_cancel.py

@@ -7,7 +7,7 @@ import asyncio
 from aioesphomeapi import ButtonInfo, CoverInfo, CoverState, EntityState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -30,9 +30,7 @@ async def test_cover_stop_tilt_cancel(
                 tilts.append(state.tilt)
 
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         cover = require_entity(entities, "test_cover", CoverInfo)
         button = require_entity(entities, "stop_tilt", ButtonInfo)

tests/integration/test_entity_missing_state.py

@@ -15,7 +15,6 @@ from aioesphomeapi import (
     ClimateMode,
     ClimateState,
     EntityInfo,
-    EntityState,
     SwitchInfo,
     SwitchState,
     WaterHeaterInfo,
@@ -24,9 +23,8 @@ from aioesphomeapi import (
 )
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
-from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -42,15 +40,6 @@ async def test_entity_missing_state(
     2. Publishing clears missing_state and reports the published value
     3. A first value that happens to equal the default still publishes
     """
-    loop = asyncio.get_running_loop()
-    futures: dict[int, asyncio.Future[EntityState]] = {}
-
-    def on_state(state: EntityState) -> None:
-        """Resolve the pending future for the entity that changed."""
-        future = futures.get(state.key)
-        if future is not None and not future.done():
-            future.set_result(state)
-
     async with (
         run_compiled(yaml_config),
         api_client_connected() as client,
@@ -76,14 +65,11 @@ async def test_entity_missing_state(
             lambda_water_heater_info,
         ]
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         # Nothing has published yet, so all of them must report unknown
         for info in stateful:
-            state = initial_state_helper.initial_states.get(info.key)
+            state = sub.initial_states.get(info.key)
             assert state is not None, f"No initial state received for {info.object_id}"
             assert state.missing_state is True, (
                 f"Initial state for {info.object_id} should have "
@@ -91,41 +77,34 @@ async def test_entity_missing_state(
             )
 
         # A configured initial_state must not start out unknown
-        initial_climate_state = initial_state_helper.initial_states.get(
-            initial_climate_info.key
-        )
+        initial_climate_state = sub.initial_states.get(initial_climate_info.key)
         assert isinstance(initial_climate_state, ClimateState)
         assert initial_climate_state.missing_state is False
         assert initial_climate_state.mode is ClimateMode.HEAT
 
         # Publishing a state on each one clears missing_state
-        futures = {info.key: loop.create_future() for info in stateful}
+        waits = {info.key: sub.waiter.next_state(info.key) for info in stateful}
         client.button_command(publish_button.key)
+        published = dict(zip(waits, await asyncio.gather(*waits.values()), strict=True))
 
-        await wait_or_fail(
-            asyncio.gather(*futures.values()),
-            5.0,
-            "Timeout waiting for published states",
-        )
-
-        switch_state = futures[switch_info.key].result()
+        switch_state = published[switch_info.key]
         assert isinstance(switch_state, SwitchState)
         assert switch_state.missing_state is False
         assert switch_state.state is True
 
-        climate_state = futures[climate_info.key].result()
+        climate_state = published[climate_info.key]
         assert isinstance(climate_state, ClimateState)
         assert climate_state.missing_state is False
         assert climate_state.mode is ClimateMode.HEAT
 
-        water_heater_state = futures[water_heater_info.key].result()
+        water_heater_state = published[water_heater_info.key]
         assert isinstance(water_heater_state, WaterHeaterState)
         assert water_heater_state.missing_state is False
         assert water_heater_state.mode is WaterHeaterMode.ECO
 
         # Its first value is OFF, which is also the default the entity starts on,
         # so it only leaves unknown if the first value publishes regardless
-        lambda_state = futures[lambda_water_heater_info.key].result()
+        lambda_state = published[lambda_water_heater_info.key]
         assert isinstance(lambda_state, WaterHeaterState)
         assert lambda_state.missing_state is False
         assert lambda_state.mode is WaterHeaterMode.OFF

tests/integration/test_fan_turn_on_action.py

@@ -6,12 +6,12 @@ field combinations and the lambda path.
 
 from __future__ import annotations
 
-import asyncio
+from functools import partial
 
-from aioesphomeapi import ButtonInfo, EntityState, FanDirection, FanInfo, FanState
+from aioesphomeapi import FanDirection, FanInfo, FanState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import press_and_wait, require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -22,54 +22,32 @@ async def test_fan_turn_on_action(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test fan TurnOnAction with constants and a lambda."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
-        fan_state_future: asyncio.Future[FanState] | None = None
-
-        def on_state(state: EntityState) -> None:
-            if (
-                isinstance(state, FanState)
-                and fan_state_future is not None
-                and not fan_state_future.done()
-            ):
-                fan_state_future.set_result(state)
-
-        async def wait_for_fan_state(timeout: float = 5.0) -> FanState:
-            nonlocal fan_state_future
-            fan_state_future = loop.create_future()
-            try:
-                return await asyncio.wait_for(fan_state_future, timeout)
-            finally:
-                fan_state_future = None
-
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         require_entity(entities, "test_fan", FanInfo)
 
-        async def press_and_wait(name: str) -> FanState:
-            btn = require_entity(entities, name.lower().replace(" ", "_"), ButtonInfo)
-            client.button_command(btn.key)
-            return await wait_for_fan_state()
+        press = partial(
+            press_and_wait, client, entities, sub.waiter, state_type=FanState
+        )
 
         # speed only
-        state = await press_and_wait("Set Speed")
+        state = await press("set_speed")
         assert state.state is True
         assert state.speed_level == 3
 
         # oscillating + direction
-        state = await press_and_wait("Set Oscillate Direction")
+        state = await press("set_oscillate_direction")
         assert state.oscillating is True
         assert state.direction == FanDirection.REVERSE
 
         # all three fields
-        state = await press_and_wait("Set All Fields")
+        state = await press("set_all_fields")
         assert state.oscillating is False
         assert state.speed_level == 4
         assert state.direction == FanDirection.FORWARD
 
         # lambda path: speed computed at runtime (test_speed global = 2)
-        state = await press_and_wait("Lambda Speed")
+        state = await press("lambda_speed")
         assert state.speed_level == 2

tests/integration/test_host_mode_climate_control.py

@@ -2,15 +2,12 @@
 
 from __future__ import annotations
 
-import asyncio
-
 import aioesphomeapi
-from aioesphomeapi import ClimateInfo, ClimateMode, EntityState
+from aioesphomeapi import ClimateInfo, ClimateMode
 import pytest
 
-from .state_utils import InitialStateHelper, require_only_entity
+from .state_utils import require_only_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
-from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -20,35 +17,25 @@ async def test_host_mode_climate_control(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test climate mode control."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
-        states: dict[int, EntityState] = {}
-        climate_future: asyncio.Future[EntityState] = loop.create_future()
-
-        def on_state(state: EntityState) -> None:
-            states[state.key] = state
-            if (
-                isinstance(state, aioesphomeapi.ClimateState)
-                and state.mode == ClimateMode.HEAT
-                and state.target_temperature_low == 21.5
-                and state.target_temperature_high == 26.5
-                and not climate_future.done()
-            ):
-                climate_future.set_result(state)
-
         # Get entities and set up state synchronization
         entities, services = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         test_climate = require_only_entity(entities, ClimateInfo)
         assert test_climate.name == "Dual-mode Thermostat"
 
-        # Subscribe with the wrapper that filters initial states
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for all initial states to be broadcast
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         # Adjust setpoints
+        climate_seen = sub.waiter.expect(
+            lambda s: (
+                isinstance(s, aioesphomeapi.ClimateState)
+                and s.mode == ClimateMode.HEAT
+                and s.target_temperature_low == 21.5
+                and s.target_temperature_high == 26.5
+            ),
+            label="Climate state",
+        )
         client.climate_command(
             test_climate.key,
             mode=ClimateMode.HEAT,
@@ -56,9 +43,7 @@ async def test_host_mode_climate_control(
             target_temperature_high=26.5,
         )
 
-        climate_state = await wait_or_fail(
-            climate_future, 5.0, "Climate state not received within 5 seconds"
-        )
+        climate_state = await climate_seen
 
         assert isinstance(climate_state, aioesphomeapi.ClimateState)
         assert climate_state.mode == ClimateMode.HEAT

tests/integration/test_light_binary_effect_off_phase.py

@@ -23,10 +23,10 @@ import asyncio
 import re
 from typing import Any
 
-from aioesphomeapi import EntityState, LightState
+from aioesphomeapi import LightState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 OUTPUT_PATTERN = re.compile(r"BINARY_OUTPUT:(YES|NO)")
@@ -66,25 +66,15 @@ async def test_light_binary_effect_off_phase(
         entities, _ = await client.list_entities_services()
         light = require_entity(entities, "test_binary_light", exact=True)
 
-        state_futures: dict[int, asyncio.Future[LightState]] = {}
-
-        def on_state(state: EntityState) -> None:
-            if isinstance(state, LightState) and state.key in state_futures:
-                future = state_futures[state.key]
-                if not future.done():
-                    future.set_result(state)
-
         # ESPHome sends the current state of every entity right after connecting; drain
         # that initial burst so it can't be mistaken for the response to a command below.
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
-        async def send_and_wait(timeout: float = 5.0, **kwargs: Any) -> LightState:
-            """Send a light command and wait for the matching state response."""
-            state_futures[light.key] = loop.create_future()
+        async def send_and_wait(**kwargs: Any) -> LightState:
+            """Send a light command and return the light's next state."""
+            next_light = sub.waiter.next_state(light.key, LightState)
             client.light_command(key=light.key, **kwargs)
-            return await asyncio.wait_for(state_futures[light.key], timeout=timeout)
+            return await next_light
 
         # A plain turn-on must drive the output on -- brightness defaults to 100% and
         # must not be mistaken for a dark phase.
@@ -165,22 +155,13 @@ async def test_light_binary_zero_brightness_is_recoverable(
         entities, _ = await client.list_entities_services()
         light = require_entity(entities, "test_binary_light", exact=True)
 
-        state_futures: dict[int, asyncio.Future[LightState]] = {}
+        sub = await subscribe_after_initial_states(client, entities)
 
-        def on_state(state: EntityState) -> None:
-            if isinstance(state, LightState) and state.key in state_futures:
-                future = state_futures[state.key]
-                if not future.done():
-                    future.set_result(state)
-
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
-
-        async def send_and_wait(timeout: float = 5.0, **kwargs: Any) -> LightState:
-            state_futures[light.key] = loop.create_future()
+        async def send_and_wait(**kwargs: Any) -> LightState:
+            """Send a light command and return the light's next state."""
+            next_light = sub.waiter.next_state(light.key, LightState)
             client.light_command(key=light.key, **kwargs)
-            return await asyncio.wait_for(state_futures[light.key], timeout=timeout)
+            return await next_light
 
         output = arm_output()
         state = await send_and_wait(state=True)

tests/integration/test_light_constant_brightness.py

@@ -12,17 +12,15 @@ gamma after the balancing formula instead of before it.
 
 from __future__ import annotations
 
-import asyncio
 import re
 from typing import Any
 
-from aioesphomeapi import EntityState, LightInfo, LightState
+from aioesphomeapi import LightInfo, LightState
 import pytest
 
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
-from .wait_utils import wait_or_fail
 
 CB_CW_RE = re.compile(r"(?<!N)CB_CW_OUTPUT:([\d.]+)")
 CB_WW_RE = re.compile(r"(?<!N)CB_WW_OUTPUT:([\d.]+)")
@@ -46,8 +44,6 @@ async def test_light_constant_brightness(
                 return float(match[1])
         return 0.0
 
-    loop = asyncio.get_running_loop()
-
     async with (
         run_compiled(yaml_config, line_callback=waiter.callback),
         api_client_connected() as client,
@@ -56,33 +52,13 @@ async def test_light_constant_brightness(
         cb_light = require_entity(entities, "test_cb_light", LightInfo, exact=True)
         ncb_light = require_entity(entities, "test_ncb_light", LightInfo, exact=True)
 
-        # Use InitialStateHelper to wait for initial state broadcast
-        initial_state_helper = InitialStateHelper(entities)
+        sub = await subscribe_after_initial_states(client, entities)
 
-        # Track state changes per light key
-        state_futures: dict[int, asyncio.Future[EntityState]] = {}
-
-        def on_state(state: EntityState) -> None:
-            if isinstance(state, LightState) and state.key in state_futures:
-                future = state_futures[state.key]
-                if not future.done():
-                    future.set_result(state)
-
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
-        await initial_state_helper.wait_for_initial_states()
-
-        async def send_and_wait(
-            light_key: int, timeout: float = 5.0, **kwargs: Any
-        ) -> LightState:
-            """Send a light command and wait for the state response."""
-            state_futures[light_key] = loop.create_future()
+        async def send_and_wait(light_key: int, **kwargs: Any) -> LightState:
+            """Send a light command and return the light's next state."""
+            next_light = sub.waiter.next_state(light_key, LightState)
             client.light_command(key=light_key, **kwargs)
-            return await wait_or_fail(
-                state_futures[light_key],
-                timeout,
-                f"Timeout waiting for light state after command: {kwargs}",
-            )
+            return await next_light
 
         # --- Test constant_brightness: true ---
 

tests/integration/test_light_control_action.py

@@ -5,12 +5,12 @@ work correctly with the compact per-field union storage. Exercises both constant
 value and lambda paths.
 """
 
-import asyncio
-from typing import Any
+from functools import partial
 
+from aioesphomeapi import LightState
 import pytest
 
-from .state_utils import require_entity
+from .state_utils import press_and_wait, require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -22,36 +22,20 @@ async def test_light_control_action(
 ) -> None:
     """Test light.turn_on, light.turn_off and light.control with constants and lambdas."""
     async with run_compiled(yaml_config), api_client_connected() as client:
-        state_futures: dict[int, asyncio.Future[Any]] = {}
-
-        def on_state(state: Any) -> None:
-            if state.key in state_futures and not state_futures[state.key].done():
-                state_futures[state.key].set_result(state)
-
-        client.subscribe_states(on_state)
-
-        # Get entities
-        entities = await client.list_entities_services()
-        light = require_entity(entities[0], "test_light", exact=True)
-        buttons = {e.name: e for e in entities[0] if hasattr(e, "name")}
-
-        async def wait_for_state(key: int, timeout: float = 5.0) -> Any:
-            """Wait for a state change for the given entity key."""
-            loop = asyncio.get_running_loop()
-            state_futures[key] = loop.create_future()
-            try:
-                return await asyncio.wait_for(state_futures[key], timeout)
-            finally:
-                state_futures.pop(key, None)
-
-        async def press_and_wait(button_name: str) -> Any:
-            """Press a button and wait for light state change."""
-            btn = buttons[button_name]
-            client.button_command(btn.key)
-            return await wait_for_state(light.key)
+        entities, _ = await client.list_entities_services()
+        light = require_entity(entities, "test_light", exact=True)
+        sub = await subscribe_after_initial_states(client, entities)
+        press = partial(
+            press_and_wait,
+            client,
+            entities,
+            sub.waiter,
+            state_type=LightState,
+            key=light.key,
+        )
 
         # Test 1: light.turn_on with RGB constants
-        state = await press_and_wait("Turn On RGB")
+        state = await press("turn_on_rgb")
         assert state.state is True
         assert state.brightness == pytest.approx(1.0)
         assert state.red == pytest.approx(0.0, abs=0.01)
@@ -59,31 +43,31 @@ async def test_light_control_action(
         assert state.blue == pytest.approx(1.0, abs=0.01)
 
         # Test 2: light.turn_off
-        state = await press_and_wait("Turn Off")
+        state = await press("turn_off")
         assert state.state is False
 
         # Test 3: light.turn_on with color_temperature
-        state = await press_and_wait("Turn On CT")
+        state = await press("turn_on_ct")
         assert state.state is True
         assert state.brightness == pytest.approx(0.8)
         assert state.color_temperature == pytest.approx(250.0)  # 4000K = 250 mireds
 
         # Test 4: light.turn_on with effect
-        state = await press_and_wait("Turn On Effect")
+        state = await press("turn_on_effect")
         assert state.effect == "Test Effect"
 
         # Test 5: Clear effect
-        state = await press_and_wait("Clear Effect")
+        state = await press("clear_effect")
         assert state.effect == "None"
 
         # Test 6: light.control with cold/warm white
-        state = await press_and_wait("Control CW")
+        state = await press("control_cw")
         assert state.cold_white == pytest.approx(0.9, abs=0.1)
         assert state.warm_white == pytest.approx(0.1, abs=0.1)
 
         # Test 7: light.turn_on with lambda brightness
         # The global test_brightness is 0.75
-        state = await press_and_wait("Lambda Brightness")
+        state = await press("lambda_brightness")
         assert state.state is True
         assert state.brightness == pytest.approx(0.75, abs=0.05)
         assert state.red == pytest.approx(1.0, abs=0.01)
@@ -91,6 +75,6 @@ async def test_light_control_action(
         assert state.blue == pytest.approx(0.0, abs=0.01)
 
         # Test 8: light.turn_on with transition_length and brightness
-        state = await press_and_wait("Turn On Transition")
+        state = await press("turn_on_transition")
         assert state.state is True
         assert state.brightness == pytest.approx(0.5)

tests/integration/test_light_dim_relative_action.py

@@ -6,12 +6,12 @@ DimRelativeAction<HasTransitionLength=true> instantiations.
 
 from __future__ import annotations
 
-import asyncio
+from functools import partial
 
-from aioesphomeapi import ButtonInfo, EntityState, LightInfo, LightState
+from aioesphomeapi import LightInfo, LightState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import press_and_wait, require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -22,51 +22,29 @@ async def test_light_dim_relative_action(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test light.dim_relative with and without transition_length."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
-        light_state_future: asyncio.Future[LightState] | None = None
-
-        def on_state(state: EntityState) -> None:
-            if (
-                isinstance(state, LightState)
-                and light_state_future is not None
-                and not light_state_future.done()
-            ):
-                light_state_future.set_result(state)
-
-        async def wait_for_light_state(timeout: float = 5.0) -> LightState:
-            nonlocal light_state_future
-            light_state_future = loop.create_future()
-            try:
-                return await asyncio.wait_for(light_state_future, timeout)
-            finally:
-                light_state_future = None
-
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         require_entity(entities, "test_light", LightInfo)
 
-        async def press_and_wait(name: str) -> LightState:
-            btn = require_entity(entities, name.lower().replace(" ", "_"), ButtonInfo)
-            client.button_command(btn.key)
-            return await wait_for_light_state()
+        press = partial(
+            press_and_wait, client, entities, sub.waiter, state_type=LightState
+        )
 
         # Setup: turn on at 50%
-        state = await press_and_wait("Setup")
+        state = await press("setup")
         assert state.state is True
         assert state.brightness == pytest.approx(0.5, abs=0.05)
 
         # Test 1: dim_relative without transition_length: 50% + 25% = 75%
-        state = await press_and_wait("Dim Up")
+        state = await press("dim_up")
         assert state.brightness == pytest.approx(0.75, abs=0.05)
 
         # Test 2: dim_relative with transition_length: 75% - 10% = 65%
-        state = await press_and_wait("Dim Down")
+        state = await press("dim_down")
         assert state.brightness == pytest.approx(0.65, abs=0.05)
 
         # Test 3: dim_relative with max_brightness limit: 65% + 50% clamped to 80%
-        state = await press_and_wait("Dim Clamp")
+        state = await press("dim_clamp")
         assert state.brightness == pytest.approx(0.80, abs=0.05)

tests/integration/test_light_effect_zero_brightness.py

@@ -17,10 +17,10 @@ import asyncio
 import re
 from typing import Any
 
-from aioesphomeapi import EntityState, LightState
+from aioesphomeapi import LightState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -46,25 +46,15 @@ async def test_light_effect_zero_brightness(
         entities, _ = await client.list_entities_services()
         light = require_entity(entities, "test_pulse_light", exact=True)
 
-        state_futures: dict[int, asyncio.Future[LightState]] = {}
-
-        def on_state(state: EntityState) -> None:
-            if isinstance(state, LightState) and state.key in state_futures:
-                future = state_futures[state.key]
-                if not future.done():
-                    future.set_result(state)
-
         # ESPHome sends the current state of every entity right after connecting; drain
         # that initial burst so it can't be mistaken for the response to a command below.
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
-        async def send_and_wait(timeout: float = 5.0, **kwargs: Any) -> LightState:
-            """Send a light command and wait for the matching state response."""
-            state_futures[light.key] = asyncio.get_running_loop().create_future()
+        async def send_and_wait(**kwargs: Any) -> LightState:
+            """Send a light command and return the light's next state."""
+            next_light = sub.waiter.next_state(light.key, LightState)
             client.light_command(key=light.key, **kwargs)
-            return await asyncio.wait_for(state_futures[light.key], timeout=timeout)
+            return await next_light
 
         # Turn the light on first so the effect starts from a known, visible state.
         state = await send_and_wait(state=True, brightness=1.0)

tests/integration/test_light_initial_state.py

@@ -7,7 +7,7 @@ that the restore logic uses as a fallback.
 
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -22,11 +22,9 @@ async def test_light_initial_state(
         entities, _ = await client.list_entities_services()
         light = require_entity(entities, "test_light")
 
-        helper = InitialStateHelper(entities)
-        client.subscribe_states(helper.on_state_wrapper(lambda s: None))
-        await helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
-        state = helper.initial_states[light.key]
+        state = sub.initial_states[light.key]
 
         # restore_mode: ALWAYS_OFF overrides state to false
         assert state.state is False
@@ -40,14 +38,14 @@ async def test_light_initial_state(
         # Regression test: RESTORE_AND_ON always forces the light on at boot, even when
         # the recovered/initial brightness was 0 -- it must never come up on-but-invisible.
         restore_and_on_light = require_entity(entities, "test_restore_and_on_light")
-        restore_and_on_state = helper.initial_states[restore_and_on_light.key]
+        restore_and_on_state = sub.initial_states[restore_and_on_light.key]
         assert restore_and_on_state.state is True
         assert restore_and_on_state.brightness == pytest.approx(1.0)
 
         # initial_state may also name the effect to start with. The light must come up
         # with that effect already active, without an on_boot automation.
         effect_light = require_entity(entities, "test_effect_light")
-        effect_state = helper.initial_states[effect_light.key]
+        effect_state = sub.initial_states[effect_light.key]
         assert effect_state.state is True
         assert effect_state.effect == "Test Pulse"
 
@@ -55,7 +53,7 @@ async def test_light_initial_state(
         # be honored as-is -- unlike every explicit restore_mode: value, omitting both
         # keys entirely must not force the light off.
         no_restore_key_light = require_entity(entities, "test_no_restore_key_light")
-        no_restore_key_state = helper.initial_states[no_restore_key_light.key]
+        no_restore_key_state = sub.initial_states[no_restore_key_light.key]
         assert no_restore_key_state.state is True
         assert no_restore_key_state.brightness == pytest.approx(0.6, abs=0.01)
 
@@ -63,7 +61,7 @@ async def test_light_initial_state(
         # light comes up red rather than the default white). RGB values are normalised to
         # proportions, so 30% red on its own is pure red.
         inferred_light = require_entity(entities, "test_inferred_color_mode_light")
-        inferred_state = helper.initial_states[inferred_light.key]
+        inferred_state = sub.initial_states[inferred_light.key]
         assert inferred_state.state is True
         assert inferred_state.red == pytest.approx(1.0, abs=0.01)
         assert inferred_state.green == pytest.approx(0.0, abs=0.01)

tests/integration/test_light_repeat_effect.py

@@ -14,10 +14,10 @@ import asyncio
 import re
 from typing import Any
 
-from aioesphomeapi import EntityState, LightState
+from aioesphomeapi import LightState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -47,23 +47,13 @@ async def test_light_repeat_effect(
         entities, _ = await client.list_entities_services()
         light = require_entity(entities, "test_rgb_light", exact=True)
 
-        state_futures: dict[int, asyncio.Future[LightState]] = {}
+        sub = await subscribe_after_initial_states(client, entities)
 
-        def on_state(state: EntityState) -> None:
-            if isinstance(state, LightState) and state.key in state_futures:
-                future = state_futures[state.key]
-                if not future.done():
-                    future.set_result(state)
-
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
-
-        async def send_and_wait(timeout: float = 5.0, **kwargs: Any) -> LightState:
-            """Send a light command and wait for the matching state response."""
-            state_futures[light.key] = asyncio.get_running_loop().create_future()
+        async def send_and_wait(**kwargs: Any) -> LightState:
+            """Send a light command and return the light's next state."""
+            next_light = sub.waiter.next_state(light.key, LightState)
             client.light_command(key=light.key, **kwargs)
-            return await asyncio.wait_for(state_futures[light.key], timeout=timeout)
+            return await next_light
 
         # Plain turn-on leaves the published colour at the white default
         state = await send_and_wait(state=True, transition_length=0.0)

tests/integration/test_light_restore_state.py

@@ -10,16 +10,14 @@ Tests that:
 
 from __future__ import annotations
 
-import asyncio
 import socket
-from typing import Any
 
 from aioesphomeapi import LightInfo, LightState
 import pytest
 
 from .conftest import run_binary_and_wait_for_port, wait_and_connect_api_client
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import CompileFunction, ConfigWriter
 
 DEVICE_NAME = "host-light-restore-state-test"
@@ -54,38 +52,26 @@ async def test_light_restore_state(
         # A single subscription serves both initial-state capture and later
         # state-change waits -- a second subscribe_states call would restart the
         # device's initial-state iterator and could resolve a wait on replayed data.
-        loop = asyncio.get_running_loop()
-        state_futures: dict[int, asyncio.Future[Any]] = {}
+        sub = await subscribe_after_initial_states(client, entities)
 
-        def on_state(state: Any) -> None:
-            if state.key in state_futures and not state_futures[state.key].done():
-                state_futures[state.key].set_result(state)
-
-        helper = InitialStateHelper(entities)
-        client.subscribe_states(helper.on_state_wrapper(on_state))
-        await helper.wait_for_initial_states()
-
-        keep_initial = helper.initial_states[keep_light.key]
+        keep_initial = sub.initial_states[keep_light.key]
         assert isinstance(keep_initial, LightState)
         assert keep_initial.state is False
 
-        override_initial = helper.initial_states[override_light.key]
+        override_initial = sub.initial_states[override_light.key]
         assert isinstance(override_initial, LightState)
         assert override_initial.state is False
 
         # Set both lights to a known, saved state
-        async def wait_for_state(key: int, timeout: float = 2.0) -> Any:
-            state_futures[key] = loop.create_future()
-            try:
-                return await asyncio.wait_for(state_futures[key], timeout)
-            finally:
-                state_futures.pop(key, None)
-
+        next_keep = sub.waiter.next_state(keep_light.key, LightState, timeout=2.0)
         client.light_command(key=keep_light.key, state=True, brightness=0.64)
-        await wait_for_state(keep_light.key)
+        await next_keep
 
+        next_override = sub.waiter.next_state(
+            override_light.key, LightState, timeout=2.0
+        )
         client.light_command(key=override_light.key, state=True, brightness=0.3)
-        await wait_for_state(override_light.key)
+        await next_override
 
     # --- Second run: same binary, same prefs file ---
     async with (
@@ -96,19 +82,17 @@ async def test_light_restore_state(
         keep_light = require_entity(entities, "test_light_keep", LightInfo)
         override_light = require_entity(entities, "test_light_override", LightInfo)
 
-        helper = InitialStateHelper(entities)
-        client.subscribe_states(helper.on_state_wrapper(lambda s: None))
-        await helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         # restore_state: {} -- comes back exactly as saved
-        keep_state = helper.initial_states[keep_light.key]
+        keep_state = sub.initial_states[keep_light.key]
         assert isinstance(keep_state, LightState)
         assert keep_state.state is True
         assert keep_state.brightness == pytest.approx(0.64, abs=0.01)
 
         # restore_state: with overrides -- state inverted, brightness forced to 100%
         # regardless of what was actually saved (0.3)
-        override_state = helper.initial_states[override_light.key]
+        override_state = sub.initial_states[override_light.key]
         assert isinstance(override_state, LightState)
         assert override_state.state is False
         assert override_state.brightness == pytest.approx(1.0, abs=0.01)

tests/integration/test_light_resume_effect.py

@@ -7,13 +7,12 @@ parameters.
 
 from __future__ import annotations
 
-import asyncio
 from typing import Any
 
-from aioesphomeapi import ColorMode, EntityState, LightInfo, LightState
+from aioesphomeapi import ColorMode, LightInfo, LightState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -32,26 +31,14 @@ async def test_light_resume_effect(
         light_resume = require_entity(entities, "test_light_resume", exact=True)
         light_no_resume = require_entity(entities, "test_light_no_resume", exact=True)
 
-        state_futures: dict[int, asyncio.Future[LightState]] = {}
-
-        def on_state(state: EntityState) -> None:
-            if isinstance(state, LightState) and state.key in state_futures:
-                future = state_futures[state.key]
-                if not future.done():
-                    future.set_result(state)
-
         # Drain initial state burst
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
-        async def send_and_wait(
-            light: LightInfo, timeout: float = 5.0, **kwargs: Any
-        ) -> LightState:
-            """Send a light command and wait for the matching state response."""
-            state_futures[light.key] = asyncio.get_running_loop().create_future()
+        async def send_and_wait(light: LightInfo, **kwargs: Any) -> LightState:
+            """Send a light command and return the light's next state."""
+            next_light = sub.waiter.next_state(light.key, LightState)
             client.light_command(key=light.key, **kwargs)
-            return await asyncio.wait_for(state_futures[light.key], timeout=timeout)
+            return await next_light
 
         # Test 1: Start with Pulse effect
         state = await send_and_wait(light_resume, state=True, effect="Pulse Effect")

tests/integration/test_light_toggle_action.py

@@ -2,12 +2,12 @@
 
 from __future__ import annotations
 
-import asyncio
+from functools import partial
 
-from aioesphomeapi import ButtonInfo, EntityState, LightInfo, LightState
+from aioesphomeapi import LightInfo, LightState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import press_and_wait, require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -18,46 +18,24 @@ async def test_light_toggle_action(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test light.toggle with and without transition_length."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
-        light_state_future: asyncio.Future[LightState] | None = None
-
-        def on_state(state: EntityState) -> None:
-            if (
-                isinstance(state, LightState)
-                and light_state_future is not None
-                and not light_state_future.done()
-            ):
-                light_state_future.set_result(state)
-
-        async def wait_for_light_state(timeout: float = 5.0) -> LightState:
-            nonlocal light_state_future
-            light_state_future = loop.create_future()
-            try:
-                return await asyncio.wait_for(light_state_future, timeout)
-            finally:
-                light_state_future = None
-
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         require_entity(entities, "test_light", LightInfo)
 
-        async def press_and_wait(name: str) -> LightState:
-            btn = require_entity(entities, name.lower().replace(" ", "_"), ButtonInfo)
-            client.button_command(btn.key)
-            return await wait_for_light_state()
+        press = partial(
+            press_and_wait, client, entities, sub.waiter, state_type=LightState
+        )
 
         # Test 1: toggle without transition_length flips off->on
-        state = await press_and_wait("Toggle")
+        state = await press("toggle")
         assert state.state is True
 
         # Test 2: toggle with transition_length flips on->off
-        state = await press_and_wait("Toggle With Trans")
+        state = await press("toggle_with_trans")
         assert state.state is False
 
         # Test 3: toggle without transition_length flips off->on again
-        state = await press_and_wait("Toggle")
+        state = await press("toggle")
         assert state.state is True

tests/integration/test_light_transition_state_publish_interval.py

@@ -16,7 +16,7 @@ from aioesphomeapi import (
 )
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity, wait_for_state
+from .state_utils import require_entity, subscribe_after_initial_states, wait_for_state
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 Timeline = list[tuple[float, LightState]]
@@ -75,10 +75,8 @@ class _Recorder:
 async def _subscribe(client: APIClient) -> tuple[list[EntityInfo], _Recorder]:
     """List entities and attach a recorder once the initial states have arrived."""
     entities, _ = await client.list_entities_services()
-    helper = InitialStateHelper(entities)
     recorder = _Recorder()
-    client.subscribe_states(helper.on_state_wrapper(recorder.on_state))
-    await helper.wait_for_initial_states()
+    await subscribe_after_initial_states(client, entities, on_state=recorder.on_state)
     return entities, recorder
 
 

tests/integration/test_preference_key_stability.py

@@ -30,7 +30,11 @@ from esphome.helpers import fnv1_hash, fnv1_hash_name, fnv1_hash_object_id
 
 from .conftest import run_binary_and_wait_for_port, wait_and_connect_api_client
 from .host_prefs import clear_host_prefs, write_host_prefs
-from .state_utils import InitialStateHelper, assert_device_name, require_entity
+from .state_utils import (
+    assert_device_name,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import CompileFunction, ConfigWriter
 
 DEVICE_NAME = "host-pref-key-stability"
@@ -102,15 +106,11 @@ async def test_preference_key_stability(
             )
             text_entity = require_entity(entities, "test_text", TextInfo, "Test Text")
 
-            initial_state_helper = InitialStateHelper(entities)
-            client.subscribe_states(
-                initial_state_helper.on_state_wrapper(lambda s: None)
-            )
-            await initial_state_helper.wait_for_initial_states()
+            sub = await subscribe_after_initial_states(client, entities)
 
-            switch_state = initial_state_helper.initial_states[switch_entity.key]
-            number_state = initial_state_helper.initial_states[number_entity.key]
-            text_state = initial_state_helper.initial_states[text_entity.key]
+            switch_state = sub.initial_states[switch_entity.key]
+            number_state = sub.initial_states[number_entity.key]
+            text_state = sub.initial_states[text_entity.key]
             assert isinstance(switch_state, SwitchState)
             assert isinstance(number_state, NumberState)
             assert isinstance(text_state, TextState)

tests/integration/test_sensor_filters_delta.py

@@ -8,7 +8,7 @@ import math
 from aioesphomeapi import ButtonInfo, EntityState, SensorState
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping
+from .state_utils import build_key_to_entity_mapping, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -109,14 +109,8 @@ async def test_sensor_filters_delta(
             },
         )
 
-        # Set up initial state helper with all entities
-        initial_state_helper = InitialStateHelper(entities)
-
-        # Subscribe to state changes with wrapper
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for initial states
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Find all buttons
         button_name_map = {

tests/integration/test_sensor_filters_rh_correction.py

@@ -5,7 +5,11 @@ from __future__ import annotations
 from aioesphomeapi import ButtonInfo
 import pytest
 
-from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
+from .state_utils import (
+    SensorStateCollector,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -29,11 +33,9 @@ async def test_sensor_filters_rh_correction(
     ):
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
         )
-        await initial_state_helper.wait_for_initial_states()
 
         done = collector.add_waiter(
             lambda: all(

tests/integration/test_sensor_filters_ring_buffer.py

@@ -7,7 +7,11 @@ import asyncio
 from aioesphomeapi import EntityState, SensorState
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping, require_entity
+from .state_utils import (
+    build_key_to_entity_mapping,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -77,14 +81,8 @@ async def test_sensor_filters_ring_buffer(
             ],
         )
 
-        # Set up initial state helper with all entities
-        initial_state_helper = InitialStateHelper(entities)
-
-        # Subscribe to state changes with wrapper
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for initial states to be sent before pressing button
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Find the publish button
         publish_button = require_entity(entities, "publish_values_button")

tests/integration/test_sensor_filters_sliding_window.py

@@ -7,7 +7,11 @@ import asyncio
 from aioesphomeapi import EntityState, SensorState
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping, require_entity
+from .state_utils import (
+    build_key_to_entity_mapping,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -109,14 +113,8 @@ async def test_sensor_filters_sliding_window(
             ],
         )
 
-        # Set up initial state helper with all entities
-        initial_state_helper = InitialStateHelper(entities)
-
-        # Subscribe to state changes with wrapper
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for initial states to be sent before pressing button
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Find the publish button
         publish_button = require_entity(entities, "publish_values_button")
@@ -231,14 +229,8 @@ async def test_sensor_filters_nan_handling(
         # Build key-to-sensor mapping
         key_to_sensor = build_key_to_entity_mapping(entities, ["min_nan", "max_nan"])
 
-        # Set up initial state helper with all entities
-        initial_state_helper = InitialStateHelper(entities)
-
-        # Subscribe to state changes with wrapper
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for initial states
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Find the publish button
         publish_button = require_entity(entities, "publish_nan_values_button")
@@ -323,14 +315,8 @@ async def test_sensor_filters_ring_buffer_wraparound(
         # Build key-to-sensor mapping
         key_to_sensor = build_key_to_entity_mapping(entities, ["wraparound_min"])
 
-        # Set up initial state helper with all entities
-        initial_state_helper = InitialStateHelper(entities)
-
-        # Subscribe to state changes with wrapper
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for initial state
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Find the publish button
         publish_button = require_entity(entities, "publish_wraparound_button")

tests/integration/test_sensor_filters_value_list.py

@@ -8,7 +8,7 @@ import math
 from aioesphomeapi import ButtonInfo, EntityState, SensorState
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping
+from .state_utils import build_key_to_entity_mapping, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -116,14 +116,8 @@ async def test_sensor_filters_value_list(
             },
         )
 
-        # Set up initial state helper with all entities
-        initial_state_helper = InitialStateHelper(entities)
-
-        # Subscribe to state changes with wrapper
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for initial states
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Find all buttons
         button_name_map = {

tests/integration/test_sensor_timeout_filter.py

@@ -7,7 +7,11 @@ import asyncio
 from aioesphomeapi import EntityState, SensorState
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping, require_entity
+from .state_utils import (
+    build_key_to_entity_mapping,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -87,10 +91,7 @@ async def test_sensor_timeout_filter(
             ],
         )
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Helper to find buttons by object_id substring
         def find_button(object_id_substring: str) -> int:

tests/integration/test_sensor_timeout_throttle_filter.py

@@ -8,7 +8,11 @@ import math
 from aioesphomeapi import ButtonInfo, EntityState, SensorState
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping, require_entity
+from .state_utils import (
+    build_key_to_entity_mapping,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -39,9 +43,7 @@ async def test_sensor_timeout_throttle_filter(
         key_to_sensor = build_key_to_entity_mapping(
             entities, {"merged": "Merged", "reference": "Reference"}
         )
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         button = require_entity(entities, "drive", ButtonInfo, exact=True)
         client.button_command(button.key)

tests/integration/test_status_flags.py

@@ -19,7 +19,11 @@ import asyncio
 
 import pytest
 
-from .state_utils import InitialStateHelper, SensorTracker, build_key_to_entity_mapping
+from .state_utils import (
+    SensorTracker,
+    build_key_to_entity_mapping,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 # Time to let the host-mode main loop run so status_led_light.loop() can
@@ -69,9 +73,9 @@ async def test_status_flags(
 
         # Swallow initial state broadcasts so the test only reacts to
         # state changes triggered by our service calls.
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(tracker.on_state))
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(
+            client, entities, on_state=tracker.on_state
+        )
 
         async def call(name: str) -> None:
             await client.execute_service(svc[name], {})

tests/integration/test_status_momentary.py

@@ -8,7 +8,7 @@ from aioesphomeapi import BinarySensorState, EntityState
 import pytest
 
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper
+from .state_utils import subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -37,11 +37,9 @@ async def test_status_momentary(
             if isinstance(state, BinarySensorState):
                 changes.put_nowait((state.key, state.state, loop.time()))
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
-        assert initial_state_helper.initial_states[warning_key].state is False
-        assert initial_state_helper.initial_states[error_key].state is False
+        sub = await subscribe_after_initial_states(client, entities, on_state=on_state)
+        assert sub.initial_states[warning_key].state is False
+        assert sub.initial_states[error_key].state is False
 
         async def next_change(key: int) -> tuple[bool, float]:
             while True:

tests/integration/test_strftime_to.py

@@ -7,7 +7,11 @@ import asyncio
 from aioesphomeapi import EntityState, TextSensorState
 import pytest
 
-from .state_utils import InitialStateHelper, assert_device_name, require_entity
+from .state_utils import (
+    assert_device_name,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -39,7 +43,7 @@ async def test_strftime_to(
             entities, "time_error_format", description="Time Error Format sensor"
         )
 
-        # Set up state tracking with InitialStateHelper
+        # Set up state tracking
         loop = asyncio.get_running_loop()
         states: dict[int, TextSensorState] = {}
         all_received = loop.create_future()
@@ -49,7 +53,6 @@ async def test_strftime_to(
             string_format.key,
             error_format.key,
         }
-        initial_state_helper = InitialStateHelper(entities)
 
         def on_state(state: EntityState) -> None:
             if isinstance(state, TextSensorState) and not state.missing_state:
@@ -57,11 +60,8 @@ async def test_strftime_to(
                 if expected_keys <= states.keys() and not all_received.done():
                     all_received.set_result(True)
 
-        # Subscribe with the wrapper that filters initial states
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for initial states to be broadcast
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Wait for all expected states
         await wait_or_fail(

tests/integration/test_template_alarm_control_panel_many_sensors.py

@@ -6,7 +6,7 @@ import aioesphomeapi
 from aioesphomeapi.model import APIIntEnum
 import pytest
 
-from .state_utils import InitialStateHelper, require_only_entity
+from .state_utils import require_only_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -66,20 +66,11 @@ async def test_template_alarm_control_panel_many_sensors(
             f"Got: {actual_sensor_names}"
         )
 
-        # Use InitialStateHelper to wait for all initial states
-        state_helper = InitialStateHelper(entities)
-
-        def on_state(state: aioesphomeapi.EntityState) -> None:
-            # We'll receive subsequent states here after initial states
-            pass
-
-        client.subscribe_states(state_helper.on_state_wrapper(on_state))
-
         # Wait for all initial states
-        await state_helper.wait_for_initial_states(timeout=5.0)
+        sub = await subscribe_after_initial_states(client, entities)
 
         # Verify the alarm state is disarmed initially
-        alarm_state = state_helper.initial_states.get(alarm_info.key)
+        alarm_state = sub.initial_states.get(alarm_info.key)
         assert alarm_state is not None, "Alarm control panel initial state not received"
         assert isinstance(alarm_state, aioesphomeapi.AlarmControlPanelEntityState)
         assert alarm_state.state == aioesphomeapi.AlarmControlPanelState.DISARMED, (
@@ -88,7 +79,7 @@ async def test_template_alarm_control_panel_many_sensors(
 
         # Verify all 10 binary sensors have initial states
         binary_sensor_states = [
-            state_helper.initial_states.get(sensor.key) for sensor in binary_sensors
+            sub.initial_states.get(sensor.key) for sensor in binary_sensors
         ]
         assert all(state is not None for state in binary_sensor_states), (
             "Not all binary sensors have initial states"

tests/integration/test_template_climate_basic.py

@@ -27,9 +27,9 @@ import pytest
 
 from .host_prefs import clear_host_prefs
 from .state_utils import (
-    InitialStateHelper,
     require_entity,
     require_only_entity,
+    subscribe_after_initial_states,
     wait_for_state,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
@@ -65,7 +65,6 @@ async def test_template_climate_basic(
             )
 
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         test_climate = require_only_entity(entities, ClimateInfo)
 
         # Advertised capabilities come straight from the supported_*/custom_* config lists.
@@ -85,12 +84,9 @@ async def test_template_climate_basic(
 
         report_button = require_entity(entities, "simulate_device_report", ButtonInfo)
 
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(lambda state: None)
-        )
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
-        initial = initial_state_helper.initial_states.get(test_climate.key)
+        initial = sub.initial_states.get(test_climate.key)
         assert initial is not None, "No initial climate state received"
         assert isinstance(initial, aioesphomeapi.ClimateState)
         assert initial.current_temperature == pytest.approx(22.5, abs=0.1)

tests/integration/test_template_climate_custom_modes.py

@@ -13,9 +13,9 @@ import pytest
 from .host_prefs import clear_host_prefs
 from .log_utils import LineWaiter
 from .state_utils import (
-    InitialStateHelper,
     require_entity,
     require_only_entity,
+    subscribe_after_initial_states,
     wait_for_state,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
@@ -47,7 +47,6 @@ async def test_template_climate_custom_modes(
             )
 
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         test_climate = require_only_entity(entities, ClimateInfo)
 
         assert set(test_climate.supported_custom_fan_modes) == {
@@ -63,12 +62,9 @@ async def test_template_climate_custom_modes(
 
         report_button = require_entity(entities, "simulate_device_report", ButtonInfo)
 
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(lambda state: None)
-        )
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
-        initial = initial_state_helper.initial_states.get(test_climate.key)
+        initial = sub.initial_states.get(test_climate.key)
         assert initial is not None, "No initial climate state received"
         assert isinstance(initial, aioesphomeapi.ClimateState)
         assert initial.custom_fan_mode == ""

tests/integration/test_template_climate_nonoptimistic.py

@@ -23,9 +23,9 @@ import pytest
 from .host_prefs import clear_host_prefs
 from .log_utils import LineWaiter
 from .state_utils import (
-    InitialStateHelper,
     require_entity,
     require_only_entity,
+    subscribe_after_initial_states,
     wait_for_state,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
@@ -55,17 +55,15 @@ async def test_template_climate_nonoptimistic(
                 state_updates.append(state)
 
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         test_climate = require_only_entity(entities, ClimateInfo)
 
         confirm_button = require_entity(
             entities, "simulate_device_confirmation", ButtonInfo
         )
 
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities, on_state=on_state)
 
-        initial = initial_state_helper.initial_states.get(test_climate.key)
+        initial = sub.initial_states.get(test_climate.key)
         assert initial is not None, "No initial climate state received"
         assert isinstance(initial, aioesphomeapi.ClimateState)
         assert initial.mode == ClimateMode.OFF

tests/integration/test_template_climate_on_control_ordering.py

@@ -13,7 +13,11 @@ import pytest
 
 from .host_prefs import clear_host_prefs
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, require_only_entity, wait_for_state
+from .state_utils import (
+    require_only_entity,
+    subscribe_after_initial_states,
+    wait_for_state,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "tmpl-clim-oc-order"
@@ -43,13 +47,9 @@ async def test_template_climate_on_control_ordering(
             )
 
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         test_climate = require_only_entity(entities, ClimateInfo)
 
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(lambda state: None)
-        )
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities)
 
         client.climate_command(test_climate.key, mode=ClimateMode.HEAT)
         state = await wait_for_climate_state()

tests/integration/test_template_climate_publish_all_fields.py

@@ -26,9 +26,9 @@ import pytest
 from .host_prefs import clear_host_prefs
 from .log_utils import LineWaiter
 from .state_utils import (
-    InitialStateHelper,
     require_entity,
     require_only_entity,
+    subscribe_after_initial_states,
     wait_for_state,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
@@ -58,13 +58,11 @@ async def test_template_climate_publish_all_fields(
                 state_updates.append(state)
 
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         require_only_entity(entities, ClimateInfo)
 
         publish_button = require_entity(entities, "publish_all", ButtonInfo)
 
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         client.button_command(publish_button.key)
         state = await wait_for_state(

tests/integration/test_template_climate_sensor_push.py

@@ -16,9 +16,9 @@ import pytest
 
 from .host_prefs import clear_host_prefs
 from .state_utils import (
-    InitialStateHelper,
     require_entity,
     require_only_entity,
+    subscribe_after_initial_states,
     wait_for_state,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
@@ -47,7 +47,6 @@ async def test_template_climate_sensor_push(
                 state_updates.append(state)
 
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         test_climate = require_only_entity(entities, ClimateInfo)
 
         publish_temp = require_entity(entities, "publish_temperature", ButtonInfo)
@@ -56,10 +55,9 @@ async def test_template_climate_sensor_push(
         )
         publish_humidity = require_entity(entities, "publish_humidity", ButtonInfo)
 
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities, on_state=on_state)
 
-        initial = initial_state_helper.initial_states.get(test_climate.key)
+        initial = sub.initial_states.get(test_climate.key)
         assert initial is not None, "No initial climate state received"
         assert isinstance(initial, aioesphomeapi.ClimateState)
         # Neither backing sensor has published anything yet.

tests/integration/test_template_climate_set_actions.py

@@ -20,9 +20,9 @@ import pytest
 from .host_prefs import clear_host_prefs
 from .log_utils import LineWaiter
 from .state_utils import (
-    InitialStateHelper,
     require_entity,
     require_only_entity,
+    subscribe_after_initial_states,
     wait_for_state,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
@@ -49,7 +49,6 @@ async def test_template_climate_set_actions(
         api_client_connected() as client,
     ):
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         test_climate = require_only_entity(entities, ClimateInfo)
 
         report_button = require_entity(entities, "report_device_state", ButtonInfo)
@@ -57,10 +56,7 @@ async def test_template_climate_set_actions(
             entities, "report_unsupported_mode", ButtonInfo
         )
 
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(lambda state: None)
-        )
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities)
 
         # Both traits are derived from the low/high and humidity set actions, not declared.
         assert test_climate.supports_two_point_target_temperature

tests/integration/test_template_climate_two_point_temperature.py

@@ -15,9 +15,9 @@ import pytest
 from .host_prefs import clear_host_prefs
 from .log_utils import LineWaiter
 from .state_utils import (
-    InitialStateHelper,
     require_entity,
     require_only_entity,
+    subscribe_after_initial_states,
     wait_for_state,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
@@ -49,7 +49,6 @@ async def test_template_climate_two_point_temperature(
             )
 
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         test_climate = require_only_entity(entities, ClimateInfo)
         assert test_climate.name == "Test Two-Point Heatpump"
         assert test_climate.supports_two_point_target_temperature
@@ -57,13 +56,9 @@ async def test_template_climate_two_point_temperature(
 
         report_button = require_entity(entities, "simulate_device_report", ButtonInfo)
 
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(lambda state: None)
-        )
+        sub = await subscribe_after_initial_states(client, entities)
 
-        await initial_state_helper.wait_for_initial_states()
-
-        initial = initial_state_helper.initial_states.get(test_climate.key)
+        initial = sub.initial_states.get(test_climate.key)
         assert initial is not None, "No initial climate state received"
         assert isinstance(initial, aioesphomeapi.ClimateState)
         # Nothing has been published yet: settable fields have no sensor to seed them from, so

tests/integration/test_template_text_save.py

@@ -10,7 +10,6 @@ from __future__ import annotations
 
 import asyncio
 import socket
-from typing import Any
 
 from aioesphomeapi import TextInfo, TextState
 import pytest
@@ -18,7 +17,11 @@ import pytest
 from .conftest import run_binary_and_wait_for_port, wait_and_connect_api_client
 from .host_prefs import clear_host_prefs
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, assert_device_name, require_entity
+from .state_utils import (
+    assert_device_name,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import CompileFunction, ConfigWriter
 
 DEVICE_NAME = "host-template-text-save-test"
@@ -60,33 +63,17 @@ async def test_template_text_save(
             entities, "test_text_restore", TextInfo, "Test Text Restore"
         )
 
-        # Set up state tracking
-        loop = asyncio.get_running_loop()
-        state_futures: dict[int, asyncio.Future[Any]] = {}
-
-        def on_state(state: Any) -> None:
-            if state.key in state_futures and not state_futures[state.key].done():
-                state_futures[state.key].set_result(state)
-
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         # Verify initial value from config
-        initial = initial_state_helper.initial_states[text_entity.key]
+        initial = sub.initial_states[text_entity.key]
         assert isinstance(initial, TextState)
         assert initial.state == "hello"
 
-        async def wait_for_state(key: int, timeout: float = 2.0) -> Any:
-            state_futures[key] = loop.create_future()
-            try:
-                return await asyncio.wait_for(state_futures[key], timeout)
-            finally:
-                state_futures.pop(key, None)
-
         # Set a new value that fits within max_length
+        next_text = sub.waiter.next_state(text_entity.key, TextState, timeout=2.0)
         client.text_command(key=text_entity.key, state="world")
-        state = await wait_for_state(text_entity.key)
+        state = await next_text
         assert state.state == "world"
 
         # Set the same value again - should NOT produce "too long" warning
@@ -112,12 +99,10 @@ async def test_template_text_save(
             entities, "test_text_restore", TextInfo, "Test Text Restore"
         )
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(lambda s: None))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         # The value should be "world" - restored from preferences
-        restored = initial_state_helper.initial_states[text_entity.key]
+        restored = sub.initial_states[text_entity.key]
         assert isinstance(restored, TextState)
         assert restored.state == "world", (
             f"Expected restored value 'world', got '{restored.state}'"

tests/integration/test_template_time.py

@@ -7,7 +7,7 @@ import asyncio
 from aioesphomeapi import EntityState, TextSensorState
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping
+from .state_utils import build_key_to_entity_mapping, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -39,9 +39,7 @@ async def test_template_time(
                 if name is not None and state.state == EXPECTED_STATES[name]:
                     events[name].set()
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         for name, event in events.items():
             await wait_or_fail(

tests/integration/test_text_command.py

@@ -7,13 +7,14 @@ transmitted via the API.
 
 from __future__ import annotations
 
-import asyncio
-from typing import Any
-
 from aioesphomeapi import TextInfo, TextState
 import pytest
 
-from .state_utils import InitialStateHelper, assert_device_name, require_entity
+from .state_utils import (
+    assert_device_name,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -24,7 +25,6 @@ async def test_text_command(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test text command with various string values including edge cases."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
         # Verify we can get device info
         await assert_device_name(client, "host-text-command-test")
@@ -41,81 +41,57 @@ async def test_text_command(
             entities, "test_text_long", TextInfo, "Test Text Long entity"
         )
 
-        # Track state changes
-        states: dict[int, Any] = {}
-        state_futures: dict[int, asyncio.Future[Any]] = {}
-
-        def on_state(state: Any) -> None:
-            states[state.key] = state
-            if state.key in state_futures and not state_futures[state.key].done():
-                state_futures[state.key].set_result(state)
-
-        # Set up InitialStateHelper to swallow initial state broadcasts
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for all initial states to be received
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         # Verify initial states were received
-        assert test_text.key in initial_state_helper.initial_states
-        initial_text_state = initial_state_helper.initial_states[test_text.key]
+        assert test_text.key in sub.initial_states
+        initial_text_state = sub.initial_states[test_text.key]
         assert isinstance(initial_text_state, TextState)
         assert initial_text_state.state == "initial"
 
-        async def wait_for_state_change(key: int, timeout: float = 2.0) -> Any:
-            """Wait for a state change for the given entity key."""
-            state_futures[key] = loop.create_future()
-            try:
-                return await asyncio.wait_for(state_futures[key], timeout)
-            finally:
-                state_futures.pop(key, None)
+        async def send_and_wait(key: int, value: str) -> TextState:
+            """Send a text command and return the entity's next state."""
+            next_text = sub.waiter.next_state(key, TextState, timeout=2.0)
+            client.text_command(key=key, state=value)
+            return await next_text
 
         # Test 1: Simple text value
-        client.text_command(key=test_text.key, state="hello world")
-        state = await wait_for_state_change(test_text.key)
+        state = await send_and_wait(test_text.key, "hello world")
         assert state.state == "hello world"
 
         # Test 2: Empty string (edge case for zero-copy)
-        client.text_command(key=test_text.key, state="")
-        state = await wait_for_state_change(test_text.key)
+        state = await send_and_wait(test_text.key, "")
         assert state.state == ""
 
         # Test 3: Single character
-        client.text_command(key=test_text.key, state="x")
-        state = await wait_for_state_change(test_text.key)
+        state = await send_and_wait(test_text.key, "x")
         assert state.state == "x"
 
         # Test 4: String with special characters
-        client.text_command(key=test_text.key, state="hello\tworld\n!")
-        state = await wait_for_state_change(test_text.key)
+        state = await send_and_wait(test_text.key, "hello\tworld\n!")
         assert state.state == "hello\tworld\n!"
 
         # Test 5: Unicode characters
-        client.text_command(key=test_text.key, state="hello 世界 🌍")
-        state = await wait_for_state_change(test_text.key)
+        state = await send_and_wait(test_text.key, "hello 世界 🌍")
         assert state.state == "hello 世界 🌍"
 
         # Test 6: Long string (tests buffer handling)
         long_text = "a" * 200
-        client.text_command(key=test_text_long.key, state=long_text)
-        state = await wait_for_state_change(test_text_long.key)
+        state = await send_and_wait(test_text_long.key, long_text)
         assert state.state == long_text
         assert len(state.state) == 200
 
         # Test 7: Password field (same mechanism, different mode)
-        client.text_command(key=test_password.key, state="newpassword123")
-        state = await wait_for_state_change(test_password.key)
+        state = await send_and_wait(test_password.key, "newpassword123")
         assert state.state == "newpassword123"
 
         # Test 8: String with null bytes embedded (edge case)
         # Note: protobuf strings should handle this but it's good to verify
-        client.text_command(key=test_text.key, state="before\x00after")
-        state = await wait_for_state_change(test_text.key)
+        state = await send_and_wait(test_text.key, "before\x00after")
         assert state.state == "before\x00after"
 
         # Test 9: Rapid successive commands (tests buffer reuse)
         for i in range(5):
-            client.text_command(key=test_text.key, state=f"rapid_{i}")
-            state = await wait_for_state_change(test_text.key)
+            state = await send_and_wait(test_text.key, f"rapid_{i}")
             assert state.state == f"rapid_{i}"

tests/integration/test_uart_mock_ld2410.py

@@ -22,7 +22,11 @@ from aioesphomeapi import ButtonInfo
 import pytest
 
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
+from .state_utils import (
+    SensorStateCollector,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -70,14 +74,10 @@ async def test_uart_mock_ld2410(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
 
-        # Set up initial state helper
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
         )
 
-        await initial_state_helper.wait_for_initial_states()
-
         # Start the UART mock scenario now that we're subscribed
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
@@ -188,13 +188,10 @@ async def test_uart_mock_ld2410_engineering(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
         )
 
-        await initial_state_helper.wait_for_initial_states()
-
         # Start the UART mock scenario now that we're subscribed
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)

tests/integration/test_uart_mock_ld2410_query_ack.py

@@ -21,7 +21,11 @@ from aioesphomeapi import (
 import pytest
 
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping, require_entity
+from .state_utils import (
+    build_key_to_entity_mapping,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -113,10 +117,7 @@ async def test_uart_mock_ld2410_query_ack(
             f"Sensor entities missing, found: {sorted(sensor_key_to_name.values())}"
         )
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         # Start the UART mock scenario now that we're subscribed
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)

tests/integration/test_uart_mock_ld2410s.py

@@ -41,7 +41,11 @@ import re
 from aioesphomeapi import APIClient, ButtonInfo, EntityState, NumberInfo, NumberState
 import pytest
 
-from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
+from .state_utils import (
+    SensorStateCollector,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -115,9 +119,7 @@ async def _start_scenario(client: APIClient, collector: SensorStateCollector) ->
     """Subscribe the collector past the initial states, then start the mock scenario."""
     entities, _ = await client.list_entities_services()
     collector.build_key_mapping(entities)
-    initial_state_helper = InitialStateHelper(entities)
-    client.subscribe_states(initial_state_helper.on_state_wrapper(collector.on_state))
-    await initial_state_helper.wait_for_initial_states()
+    await subscribe_after_initial_states(client, entities, on_state=collector.on_state)
     start_btn = require_entity(entities, "start_scenario", ButtonInfo)
     client.button_command(start_btn.key)
 
@@ -194,12 +196,10 @@ async def test_uart_mock_ld2410s(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
         max_distance_key = require_entity(entities, "max_distance", NumberInfo).key
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities, on_state=on_state)
         # The values read from the module: gate 16 and 10 s
         no_delay_key = require_entity(entities, "no_delay", NumberInfo).key
-        initial = initial_state_helper.initial_states
+        initial = sub.initial_states
         assert initial[max_distance_key].state == pytest.approx(11.2)
         assert initial[no_delay_key].state == pytest.approx(10)
 

tests/integration/test_uart_mock_ld2412.py

@@ -32,22 +32,15 @@ import asyncio
 from collections.abc import Callable
 import math
 
-from aioesphomeapi import (
-    APIClient,
-    ButtonInfo,
-    EntityInfo,
-    EntityState,
-    NumberInfo,
-    NumberState,
-)
+from aioesphomeapi import ButtonInfo, EntityInfo, EntityState, NumberInfo, NumberState
 import pytest
 
 from .log_utils import LineWaiter
 from .state_utils import (
-    InitialStateHelper,
     SensorStateCollector,
     StateWaiter,
     require_entity,
+    subscribe_after_initial_states,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
@@ -82,18 +75,6 @@ def _threshold_numbers(entities: list[EntityInfo]) -> dict[str, NumberInfo]:
     }
 
 
-async def _wait_initial(
-    client: APIClient,
-    entities: list[EntityInfo],
-    on_state: Callable[[EntityState], None] = lambda s: None,
-) -> InitialStateHelper:
-    """Subscribe to states and wait until every entity has sent its initial state."""
-    initial_state_helper = InitialStateHelper(entities)
-    client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-    await initial_state_helper.wait_for_initial_states()
-    return initial_state_helper
-
-
 def _number_state_of(key: int) -> Callable[[EntityState], bool]:
     """Return a predicate that matches a number state with a value for ``key``."""
     return lambda s: isinstance(s, NumberState) and not s.missing_state and s.key == key
@@ -175,7 +156,9 @@ async def test_uart_mock_ld2412(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
 
-        await _wait_initial(client, entities, collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
+        )
 
         # Start the UART mock scenario now that we're subscribed
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
@@ -275,7 +258,9 @@ async def test_uart_mock_ld2412_engineering(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
 
-        await _wait_initial(client, entities, collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
+        )
 
         # Start the UART mock scenario now that we're subscribed
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
@@ -362,7 +347,9 @@ async def test_uart_mock_ld2412_engineering_truncated(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
 
-        await _wait_initial(client, entities, collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
+        )
 
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
@@ -423,7 +410,7 @@ async def test_uart_mock_ld2412_footer_in_payload(
         api_client_connected() as client,
     ):
         entities, _ = await client.list_entities_services()
-        initial_state_helper = await _wait_initial(client, entities)
+        sub = await subscribe_after_initial_states(client, entities)
         numbers = _threshold_numbers(entities)
 
         # Gate 5 move is the first footer byte inside the motion answer
@@ -433,7 +420,7 @@ async def test_uart_mock_ld2412_footer_in_payload(
             "gate_0_still_threshold": 0.0,
             "gate_5_still_threshold": 25.0,
         }.items():
-            state = initial_state_helper.initial_states[numbers[name].key]
+            state = sub.initial_states[numbers[name].key]
             assert state.state == pytest.approx(value), name
 
         motion_written = events.next_line_with(expected_motion)
@@ -523,10 +510,10 @@ async def test_uart_mock_ld2412_truncated_answer(
                 states.setdefault(state.key, []).append(state.state)
             state_waiter.on_state(state)
 
-        initial_state_helper = await _wait_initial(client, entities, record)
+        sub = await subscribe_after_initial_states(client, entities, on_state=record)
 
         # The late answer may already have been parsed when the client connected
-        for state in initial_state_helper.initial_states.values():
+        for state in sub.initial_states.values():
             record(state)
 
         try:
@@ -573,17 +560,16 @@ async def test_uart_mock_ld2412_requery_after_giveup(
         gate_5_key = numbers["gate_5_move_threshold"].key
 
         gate_5_values: list[float] = []
-        waiter = StateWaiter()
 
         def on_state(state: EntityState) -> None:
             if _number_state_of(gate_5_key)(state):
                 gate_5_values.append(state.state)
-            waiter.on_state(state)
 
-        initial_state_helper = await _wait_initial(client, entities, on_state)
+        sub = await subscribe_after_initial_states(client, entities, on_state=on_state)
+        waiter = sub.waiter
 
         for key in (gate_0_key, gate_5_key):
-            assert initial_state_helper.initial_states[key].missing_state
+            assert sub.initial_states[key].missing_state
 
         held_back = events.next_line_with(MOTION_HELD_BACK)
         config_mode_left = events.next_line_with(CONFIG_MODE_LEFT)

tests/integration/test_uart_mock_ld2420.py

@@ -23,7 +23,11 @@ from aioesphomeapi import ButtonInfo
 import pytest
 
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
+from .state_utils import (
+    SensorStateCollector,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -54,14 +58,10 @@ async def test_uart_mock_ld2420(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
 
-        # Set up initial state helper
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
         )
 
-        await initial_state_helper.wait_for_initial_states()
-
         # Start the UART mock scenario now that we're subscribed
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
@@ -151,13 +151,10 @@ async def test_uart_mock_ld2420_simple(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
         )
 
-        await initial_state_helper.wait_for_initial_states()
-
         # Start the UART mock scenario now that we're subscribed
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)

tests/integration/test_uart_mock_ld2420_config_exit.py

@@ -8,7 +8,11 @@ from aioesphomeapi import ButtonInfo
 import pytest
 
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
+from .state_utils import (
+    SensorStateCollector,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -47,11 +51,9 @@ async def test_uart_mock_ld2420_config_exit(
     ):
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
         )
-        await initial_state_helper.wait_for_initial_states()
 
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)

tests/integration/test_uart_mock_ld2450.py

@@ -21,7 +21,11 @@ from aioesphomeapi import ButtonInfo
 import pytest
 
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
+from .state_utils import (
+    SensorStateCollector,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -83,14 +87,10 @@ async def test_uart_mock_ld2450(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
 
-        # Set up initial state helper
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
         )
 
-        await initial_state_helper.wait_for_initial_states()
-
         # Start the UART mock scenario now that we're subscribed
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)

tests/integration/test_uart_mock_ld2450_has_target_only.py

@@ -12,7 +12,7 @@ import asyncio
 from aioesphomeapi import BinarySensorInfo, BinarySensorState, ButtonInfo, EntityState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 # The fixture sends one frame with a target, then 2 s of frames without one
@@ -41,9 +41,7 @@ async def test_uart_mock_ld2450_has_target_only(
                 if state.state and not turned_on.done():
                     turned_on.set_result(None)
 
-        helper = InitialStateHelper(entities)
-        client.subscribe_states(helper.on_state_wrapper(on_state))
-        await helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)

tests/integration/test_uart_mock_ld2450_polygon_zone_restore.py

@@ -7,7 +7,6 @@ Tests:
 
 from __future__ import annotations
 
-import asyncio
 import socket
 import struct
 
@@ -17,7 +16,7 @@ import pytest
 from .conftest import run_binary_and_wait_for_port, wait_and_connect_api_client
 from .host_prefs import clear_host_prefs, read_host_prefs, write_host_prefs
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import CompileFunction, ConfigWriter
 
 DEVICE_NAME = "uart-mock-ld2450-zone-restore"
@@ -53,10 +52,8 @@ async def test_uart_mock_ld2450_polygon_zone_restore(
         ):
             entities, _ = await client.list_entities_services()
             polygon = require_entity(entities, "restore_zone_polygon", TextInfo)
-            helper = InitialStateHelper(entities)
-            client.subscribe_states(helper.on_state_wrapper(lambda state: None))
-            await helper.wait_for_initial_states()
-            state = helper.initial_states[polygon.key]
+            sub = await subscribe_after_initial_states(client, entities)
+            state = sub.initial_states[polygon.key]
             assert isinstance(state, TextState)
             return state.state
 
@@ -67,25 +64,18 @@ async def test_uart_mock_ld2450_polygon_zone_restore(
     ):
         entities, _ = await client.list_entities_services()
         polygon = require_entity(entities, "restore_zone_polygon", TextInfo)
-        loop = asyncio.get_running_loop()
-        published = loop.create_future()
-
-        def on_state(state: object) -> None:
-            if (
-                isinstance(state, TextState)
-                and state.key == polygon.key
-                and state.state == ZONE_POLYGON
-                and not published.done()
-            ):
-                published.set_result(None)
-
-        helper = InitialStateHelper(entities)
-        client.subscribe_states(helper.on_state_wrapper(on_state))
-        await helper.wait_for_initial_states()
-        assert helper.initial_states[polygon.key].state == ""
+        sub = await subscribe_after_initial_states(client, entities)
+        assert sub.initial_states[polygon.key].state == ""
 
+        published = sub.waiter.expect(
+            lambda s: (
+                isinstance(s, TextState)
+                and s.key == polygon.key
+                and s.state == ZONE_POLYGON
+            )
+        )
         client.text_command(polygon.key, ZONE_POLYGON)
-        await asyncio.wait_for(published, timeout=5.0)
+        await published
 
     # Second run: the polygon is restored
     assert await initial_polygon() == ZONE_POLYGON

tests/integration/test_uart_mock_ld2450_polygon_zone_timeout.py

@@ -22,7 +22,7 @@ from aioesphomeapi import (
 import pytest
 
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "uart-mock-ld2450-zone-timeout"
@@ -76,9 +76,7 @@ async def test_uart_mock_ld2450_polygon_zone_timeout(
             ):
                 timeout_set.set_result(None)
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        await subscribe_after_initial_states(client, entities, on_state=on_state)
 
         client.text_command(polygon.key, ZONE_POLYGON)
         client.number_command(timeout.key, PRESENCE_TIMEOUT_S)

tests/integration/test_uart_mock_ld2450_polygon_zones.py

@@ -28,7 +28,7 @@ import pytest
 from esphome.components.ld2450.text import POLYGON_MAX_LENGTH, POLYGON_PATTERN
 
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, StateWaiter, require_entity
+from .state_utils import require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "uart-mock-ld2450-zones"
@@ -64,28 +64,22 @@ async def test_uart_mock_ld2450_polygon_zones(
         presence_history: dict[str, list[bool | None]] = {
             name: [] for name in presences
         }
-        waiter = StateWaiter()
-        expect = waiter.expect
 
         def on_state(state: EntityState) -> None:
             if isinstance(state, BinarySensorState) and state.key in presence_names:
                 presence_history[presence_names[state.key]].append(
                     None if state.missing_state else state.state
                 )
-            waiter.on_state(state)
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities, on_state=on_state)
+        expect = sub.waiter.expect
 
         for name, polygon in polygons.items():
             # Home Assistant validates input and the zone card finds zones with these
             assert polygon.pattern == POLYGON_PATTERN
             assert polygon.max_length == POLYGON_MAX_LENGTH
-            assert initial_state_helper.initial_states[polygon.key].state == ""
-            assert initial_state_helper.initial_states[
-                presences[name].key
-            ].missing_state
+            assert sub.initial_states[polygon.key].state == ""
+            assert sub.initial_states[presences[name].key].missing_state
 
         async def set_polygon(name: str, value: str, expected: str) -> None:
             key = polygons[name].key

tests/integration/test_uart_mock_ld2460.py

@@ -18,7 +18,11 @@ from aioesphomeapi import ButtonInfo
 import pytest
 
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
+from .state_utils import (
+    SensorStateCollector,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -72,11 +76,9 @@ async def test_uart_mock_ld2460(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
         )
-        await initial_state_helper.wait_for_initial_states()
 
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)

tests/integration/test_uart_mock_pylontech.py

@@ -17,7 +17,11 @@ from aioesphomeapi import ButtonInfo
 import pytest
 
 from .log_utils import LineWaiter
-from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
+from .state_utils import (
+    SensorStateCollector,
+    require_entity,
+    subscribe_after_initial_states,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -88,11 +92,9 @@ async def test_uart_mock_pylontech(
         entities, _ = await client.list_entities_services()
         collector.build_key_mapping(entities)
 
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(
-            initial_state_helper.on_state_wrapper(collector.on_state)
+        await subscribe_after_initial_states(
+            client, entities, on_state=collector.on_state
         )
-        await initial_state_helper.wait_for_initial_states()
 
         start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)

tests/integration/test_valve_control_action.py

@@ -6,12 +6,12 @@ field combinations and the lambda path.
 
 from __future__ import annotations
 
-import asyncio
+from functools import partial
 
-from aioesphomeapi import ButtonInfo, EntityState, ValveInfo, ValveOperation, ValveState
+from aioesphomeapi import ValveInfo, ValveOperation, ValveState
 import pytest
 
-from .state_utils import InitialStateHelper, require_entity
+from .state_utils import press_and_wait, require_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -22,51 +22,29 @@ async def test_valve_control_action(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test valve ControlAction with constants and a lambda."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
-        valve_state_future: asyncio.Future[ValveState] | None = None
-
-        def on_state(state: EntityState) -> None:
-            if (
-                isinstance(state, ValveState)
-                and valve_state_future is not None
-                and not valve_state_future.done()
-            ):
-                valve_state_future.set_result(state)
-
-        async def wait_for_valve_state(timeout: float = 5.0) -> ValveState:
-            nonlocal valve_state_future
-            valve_state_future = loop.create_future()
-            try:
-                return await asyncio.wait_for(valve_state_future, timeout)
-            finally:
-                valve_state_future = None
-
         entities, _ = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         require_entity(entities, "test_valve", ValveInfo)
 
-        async def press_and_wait(name: str) -> ValveState:
-            btn = require_entity(entities, name.lower().replace(" ", "_"), ButtonInfo)
-            client.button_command(btn.key)
-            return await wait_for_valve_state()
+        press = partial(
+            press_and_wait, client, entities, sub.waiter, state_type=ValveState
+        )
 
         # valve.control: position only
-        state = await press_and_wait("Set Position")
+        state = await press("set_position")
         assert state.position == pytest.approx(0.5, abs=0.01)
 
         # valve.control: state alias for position 1.0
-        state = await press_and_wait("Open State")
+        state = await press("open_state")
         assert state.position == pytest.approx(1.0, abs=0.01)
 
         # valve.control: lambda position (test_position global = 0.42)
-        state = await press_and_wait("Lambda Position")
+        state = await press("lambda_position")
         assert state.position == pytest.approx(0.42, abs=0.01)
 
         # valve.control: stop only — template valve's stop_action publishes
         # current_operation: IDLE.
-        state = await press_and_wait("Stop Valve")
+        state = await press("stop_valve")
         assert state.current_operation == ValveOperation.IDLE

tests/integration/test_water_heater_template.py

@@ -14,7 +14,7 @@ from aioesphomeapi import (
 )
 import pytest
 
-from .state_utils import InitialStateHelper, require_only_entity
+from .state_utils import require_only_entity, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -25,32 +25,9 @@ async def test_water_heater_template(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test template water heater basic state and mode changes."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
-        states: dict[int, aioesphomeapi.EntityState] = {}
-        state_future: asyncio.Future[WaterHeaterState] | None = None
-
-        def on_state(state: aioesphomeapi.EntityState) -> None:
-            states[state.key] = state
-            if (
-                isinstance(state, WaterHeaterState)
-                and state_future is not None
-                and not state_future.done()
-            ):
-                state_future.set_result(state)
-
-        async def wait_for_state(timeout: float = 5.0) -> WaterHeaterState:
-            """Wait for next water heater state change."""
-            nonlocal state_future
-            state_future = loop.create_future()
-            try:
-                return await asyncio.wait_for(state_future, timeout)
-            finally:
-                state_future = None
-
         # Get entities and set up state synchronization
         entities, services = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
         test_water_heater = require_only_entity(entities, WaterHeaterInfo)
 
         # Verify water heater entity info
@@ -72,14 +49,11 @@ async def test_water_heater_template(
             f"Expected 4 supported modes, got {len(supported_modes)}: {supported_modes}"
         )
 
-        # Subscribe with the wrapper that filters initial states
-        client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
-
         # Wait for all initial states to be broadcast
-        await initial_state_helper.wait_for_initial_states()
+        sub = await subscribe_after_initial_states(client, entities)
 
         # Get initial state and verify
-        initial_state = initial_state_helper.initial_states.get(test_water_heater.key)
+        initial_state = sub.initial_states.get(test_water_heater.key)
         assert initial_state is not None, "Water heater initial state not found"
         assert isinstance(initial_state, WaterHeaterState)
         # Initial mode is OFF (default) since we don't have a mode lambda
@@ -115,37 +89,43 @@ async def test_water_heater_template(
         )
 
         # Test turning on away mode
+        next_state = sub.waiter.next_state(test_water_heater.key, WaterHeaterState)
         client.water_heater_command(test_water_heater.key, away=True)
-        away_on_state = await wait_for_state()
+        away_on_state = await next_state
         assert (away_on_state.state & WaterHeaterStateFlag.AWAY) != 0
         # ON flag should still be set (is_on lambda returns true)
         assert (away_on_state.state & WaterHeaterStateFlag.ON) != 0
 
         # Test turning off away mode
+        next_state = sub.waiter.next_state(test_water_heater.key, WaterHeaterState)
         client.water_heater_command(test_water_heater.key, away=False)
-        away_off_state = await wait_for_state()
+        away_off_state = await next_state
         assert (away_off_state.state & WaterHeaterStateFlag.AWAY) == 0
         assert (away_off_state.state & WaterHeaterStateFlag.ON) != 0
 
         # Test turning off (on=False)
+        next_state = sub.waiter.next_state(test_water_heater.key, WaterHeaterState)
         client.water_heater_command(test_water_heater.key, on=False)
-        off_state = await wait_for_state()
+        off_state = await next_state
         assert (off_state.state & WaterHeaterStateFlag.ON) == 0
         assert (off_state.state & WaterHeaterStateFlag.AWAY) == 0
 
         # Test turning back on (on=True)
+        next_state = sub.waiter.next_state(test_water_heater.key, WaterHeaterState)
         client.water_heater_command(test_water_heater.key, on=True)
-        on_state = await wait_for_state()
+        on_state = await next_state
         assert (on_state.state & WaterHeaterStateFlag.ON) != 0
 
         # Test changing to GAS mode
+        next_state = sub.waiter.next_state(test_water_heater.key, WaterHeaterState)
         client.water_heater_command(test_water_heater.key, mode=WaterHeaterMode.GAS)
-        gas_state = await wait_for_state()
+        gas_state = await next_state
         assert gas_state.mode == WaterHeaterMode.GAS
 
         # Test changing to ECO mode (from GAS)
+        next_state = sub.waiter.next_state(test_water_heater.key, WaterHeaterState)
         client.water_heater_command(test_water_heater.key, mode=WaterHeaterMode.ECO)
-        eco_state = await wait_for_state()
+        eco_state = await next_state
         assert eco_state.mode == WaterHeaterMode.ECO