Skip to content
rachel/workerd
File

Blob: src/workerd/api/streams/writable-sink-test.c++

17.3 KB
1#include "writable-sink.h"
2 
3#include <workerd/jsg/jsg-test.h>
4#include <workerd/tests/test-fixture.h>
5#include <workerd/util/own-util.h>
6#include <workerd/util/stream-utils.h>
7 
8#include <kj/async-io.h>
9#include <kj/compat/gzip.h>
10#include <kj/test.h>
11 
12namespace workerd::api::streams {
13namespace {
14 
15// Mock WritableSink for testing wrapper functionality
16class MockWritableSink final: public WritableSink {
17 public:
18  MockWritableSink() = default;
19  ~MockWritableSink() = default;
20 
21  kj::Promise<void> write(kj::ArrayPtr<const kj::byte> buffer) override {
22    writeCallCount++;
23    lastWriteSize = buffer.size();
24    totalBytesWritten += buffer.size();
25 
26    if (shouldFailWrite) {
27      KJ_FAIL_REQUIRE("Expected failure");
28    }
29 
30    // Copy data for verification
31    writtenData.addAll(buffer);
32    co_return;
33  }
34 
35  kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const kj::byte>> pieces) override {
36    multiWriteCallCount++;
37    size_t totalSize = 0;
38 
39    for (auto piece: pieces) {
40      totalSize += piece.size();
41      writtenData.addAll(piece);
42    }
43 
44    totalBytesWritten += totalSize;
45    lastWriteSize = totalSize;
46 
47    if (shouldFailWrite) {
48      KJ_FAIL_REQUIRE("Expected failure");
49    }
50 
51    co_return;
52  }
53 
54  kj::Promise<void> end() override {
55    endCallCount++;
56    isEnded = true;
57 
58    if (shouldFailEnd) {
59      KJ_FAIL_REQUIRE("Expected failure");
60    }
61 
62    co_return;
63  }
64 
65  void abort(kj::Exception reason) override {
66    abortCallCount++;
67    abortReason = kj::mv(reason);
68  }
69 
70  rpc::StreamEncoding disownEncodingResponsibility() override {
71    disownCallCount++;
72    auto prev = encoding;
73    encoding = rpc::StreamEncoding::IDENTITY;
74    return prev;
75  }
76 
77  rpc::StreamEncoding getEncoding() override {
78    getEncodingCallCount++;
79    return encoding;
80  }
81 
82  // Test state accessors
83  uint32_t writeCallCount = 0;
84  uint32_t multiWriteCallCount = 0;
85  uint32_t endCallCount = 0;
86  uint32_t abortCallCount = 0;
87  uint32_t disownCallCount = 0;
88  uint32_t getEncodingCallCount = 0;
89 
90  size_t lastWriteSize = 0;
91  size_t totalBytesWritten = 0;
92  bool isEnded = false;
93  kj::Maybe<kj::Exception> abortReason;
94 
95  kj::Vector<kj::byte> writtenData;
96  rpc::StreamEncoding encoding = rpc::StreamEncoding::IDENTITY;
97 
98  // Control behavior for testing
99  bool shouldFailWrite = false;
100  bool shouldFailEnd = false;
101};
102 
103// Test memory-based AsyncOutputStream for factory function tests
104class MemoryAsyncOutputStream final: public kj::AsyncOutputStream {
105 public:
106  kj::Promise<void> write(kj::ArrayPtr<const kj::byte> buffer) override {
107    data.addAll(buffer);
108 
109    if (writeShouldError) {
110      KJ_FAIL_REQUIRE("Expected failure");
111    }
112 
113    co_return;
114  }
115 
116  kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const kj::byte>> pieces) override {
117    for (auto piece: pieces) {
118      data.addAll(piece);
119    }
120 
121    if (writeShouldError) {
122      KJ_FAIL_REQUIRE("Expected failure");
123    }
124 
125    co_return;
126  }
127 
128  kj::Promise<void> whenWriteDisconnected() override {
129    return kj::NEVER_DONE;
130  }
131 
132  bool writeShouldError = false;
133 
134  kj::Vector<kj::byte> data;
135};
136 
137struct MockEndable final: public EndableAsyncOutputStream {
138  bool isEnded = false;
139  kj::Vector<kj::byte> data;
140 
141  kj::Promise<void> write(kj::ArrayPtr<const kj::byte> buffer) override {
142    data.addAll(buffer);
143    co_return;
144  }
145 
146  kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const kj::byte>> pieces) override {
147    for (auto piece: pieces) {
148      data.addAll(piece);
149    }
150    co_return;
151  }
152 
153  kj::Promise<void> whenWriteDisconnected() override {
154    return kj::NEVER_DONE;
155  }
156 
157  kj::Promise<void> end() override {
158    isEnded = true;
159    co_return;
160  }
161};
162 
163// ======================================================================================
164// Core WritableSink Interface Tests
165 
166KJ_TEST("WritableSink basic write operations") {
167  TestFixture fixture;
168  MockWritableSink sink;
169  kj::byte testData[] = {1, 2, 3, 4, 5};
170 
171  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
172    // Test single buffer write
173    kj::ArrayPtr<const kj::byte> buffer(testData, 5);
174 
175    co_await sink.write(buffer);
176    co_await sink.end();
177  });
178 
179  KJ_ASSERT(sink.writeCallCount == 1);
180  KJ_ASSERT(sink.lastWriteSize == 5);
181  KJ_ASSERT(sink.totalBytesWritten == 5);
182  KJ_ASSERT(sink.writtenData.size() == 5);
183  KJ_ASSERT(sink.writtenData == testData);
184  KJ_ASSERT(sink.isEnded);
185}
186 
187KJ_TEST("WritableSink multi-piece write operations") {
188  TestFixture fixture;
189  MockWritableSink sink;
190 
191  // Test multi-piece write
192  kj::byte data1[] = {1, 2, 3};
193  kj::byte data2[] = {4, 5};
194  kj::byte data3[] = {6, 7, 8, 9};
195 
196  kj::ArrayPtr<const kj::byte> pieces[] = {kj::ArrayPtr<const kj::byte>(data1, 3),
197    kj::ArrayPtr<const kj::byte>(data2, 2), kj::ArrayPtr<const kj::byte>(data3, 4)};
198 
199  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
200    co_await sink.write(kj::ArrayPtr<const kj::ArrayPtr<const kj::byte>>(pieces, 3));
201  });
202 
203  KJ_ASSERT(sink.multiWriteCallCount == 1);
204  KJ_ASSERT(sink.lastWriteSize == 9);
205  KJ_ASSERT(sink.totalBytesWritten == 9);
206  KJ_ASSERT(sink.writtenData.size() == 9);
207 
208  // Verify data order
209  kj::byte expected[] = {1, 2, 3, 4, 5, 6, 7, 8, 9};
210  KJ_ASSERT(sink.writtenData == kj::ArrayPtr<const kj::byte>(expected, 9));
211}
212 
213KJ_TEST("WritableSink end operation") {
214  TestFixture fixture;
215  MockWritableSink sink;
216 
217  fixture.runInIoContext(
218      [&](const auto& environment) -> kj::Promise<void> { co_await sink.end(); });
219 
220  KJ_ASSERT(sink.endCallCount == 1);
221  KJ_ASSERT(sink.isEnded);
222}
223 
224KJ_TEST("WritableSink abort operation") {
225  MockWritableSink sink;
226 
227  sink.abort(KJ_EXCEPTION(DISCONNECTED, "Abort reason"));
228 
229  KJ_ASSERT(sink.abortCallCount == 1);
230  KJ_ASSERT(sink.abortReason != kj::none);
231}
232 
233KJ_TEST("WritableSink encoding operations") {
234  MockWritableSink sink;
235 
236  auto encoding = sink.getEncoding();
237  KJ_ASSERT(encoding == rpc::StreamEncoding::IDENTITY);
238  KJ_ASSERT(sink.getEncodingCallCount == 1);
239 
240  auto disownedEncoding = sink.disownEncodingResponsibility();
241  KJ_ASSERT(disownedEncoding == rpc::StreamEncoding::IDENTITY);
242  KJ_ASSERT(sink.disownCallCount == 1);
243}
244 
245// ======================================================================================
246// WritableSinkWrapper Tests
247 
248KJ_TEST("WritableSinkWrapper write/end delegation") {
249  TestFixture fixture;
250  auto innerSink = kj::heap<MockWritableSink>();
251  auto& sink = *innerSink;
252  class TestWrapper: public WritableSinkWrapper {
253   public:
254    TestWrapper(kj::Own<WritableSink> inner): WritableSinkWrapper(kj::mv(inner)) {}
255  };
256  auto wrapper = kj::heap<TestWrapper>(kj::mv(innerSink));
257  kj::byte testData[] = {1, 2, 3};
258 
259  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
260    co_await wrapper->write(kj::ArrayPtr<const kj::byte>(testData, 3));
261    co_await wrapper->end();
262  });
263 
264  KJ_ASSERT(sink.writeCallCount == 1);
265  KJ_ASSERT(sink.lastWriteSize == 3);
266  KJ_ASSERT(sink.totalBytesWritten == 3);
267  KJ_ASSERT(sink.writtenData.size() == 3);
268  KJ_ASSERT(sink.writtenData == kj::arrayPtr(testData, sizeof(testData)));
269  KJ_ASSERT(sink.endCallCount == 1);
270  KJ_ASSERT(sink.isEnded);
271}
272 
273KJ_TEST("WritableSinkWrapper abort delegation") {
274  TestFixture fixture;
275  auto innerSink = kj::heap<MockWritableSink>();
276  auto& sink = *innerSink;
277  class TestWrapper: public WritableSinkWrapper {
278   public:
279    TestWrapper(kj::Own<WritableSink> inner): WritableSinkWrapper(kj::mv(inner)) {}
280  };
281  auto wrapper = kj::heap<TestWrapper>(kj::mv(innerSink));
282  wrapper->abort(KJ_EXCEPTION(FAILED, "test abort"));
283  KJ_ASSERT(sink.abortCallCount == 1);
284  KJ_ASSERT(sink.abortReason != kj::none);
285}
286 
287KJ_TEST("WritableSinkWrapper encoding delegation") {
288  TestFixture fixture;
289  auto innerSink = kj::heap<MockWritableSink>();
290  auto& sink = *innerSink;
291  sink.encoding = rpc::StreamEncoding::GZIP;
292  class TestWrapper: public WritableSinkWrapper {
293   public:
294    TestWrapper(kj::Own<WritableSink> inner): WritableSinkWrapper(kj::mv(inner)) {}
295  };
296  auto wrapper = kj::heap<TestWrapper>(kj::mv(innerSink));
297 
298  auto encoding = wrapper->getEncoding();
299  KJ_ASSERT(encoding == rpc::StreamEncoding::GZIP);
300  KJ_ASSERT(sink.getEncodingCallCount == 1);
301  auto disowned = wrapper->disownEncodingResponsibility();
302  KJ_ASSERT(disowned == rpc::StreamEncoding::GZIP);
303  KJ_ASSERT(wrapper->getEncoding() == rpc::StreamEncoding::IDENTITY);
304  KJ_ASSERT(sink.disownCallCount == 1);
305}
306 
307KJ_TEST("WritableSinkWrapper release functionality") {
308  auto innerSink = kj::heap<MockWritableSink>();
309  auto innerPtr = innerSink.get();
310 
311  class TestWrapper: public WritableSinkWrapper {
312   public:
313    TestWrapper(kj::Own<WritableSink> inner): WritableSinkWrapper(kj::mv(inner)) {}
314  };
315 
316  auto wrapper = kj::heap<TestWrapper>(kj::mv(innerSink));
317 
318  // Release the inner sink
319  auto released = wrapper->release();
320  KJ_ASSERT(released.get() == innerPtr);
321 
322  // Wrapper should no longer be usable
323  try {
324    wrapper->abort(KJ_EXCEPTION(FAILED, "test"));
325    KJ_FAIL_REQUIRE("Expected exception on using released wrapper");
326  } catch (...) {
327    auto exception = kj::getCaughtExceptionAsKj();
328    KJ_EXPECT(exception.getDescription().contains("inner != nullptr"));
329  }
330}
331 
332// ======================================================================================
333// Factory Function Tests
334 
335KJ_TEST("newWritableSink with AsyncOutputStream") {
336  TestFixture fixture;
337  MemoryAsyncOutputStream inner;
338  auto fakeOwn = kj::Own<MemoryAsyncOutputStream>(&inner, kj::NullDisposer::instance);
339  auto sink = newWritableSink(kj::mv(fakeOwn));
340  kj::byte testData[] = {1, 2, 3, 4, 5};
341 
342  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
343    // Test writing data
344    co_await sink->write(kj::ArrayPtr<const kj::byte>(testData, 5));
345    co_await sink->end();
346  });
347 
348  KJ_ASSERT(inner.data.size() == 5);
349  KJ_ASSERT(inner.data == kj::arrayPtr(testData, sizeof(testData)));
350}
351 
352KJ_TEST("newClosedWritableSink (write)") {
353  TestFixture fixture;
354  auto sink = newClosedWritableSink();
355  kj::byte testData[] = {1, 2, 3};
356 
357  try {
358    fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
359      // Test write on closed sink
360      co_await sink->write(kj::ArrayPtr<const kj::byte>(testData, 3));
361    });
362    KJ_FAIL_REQUIRE("should have failed");
363  } catch (...) {
364    auto ex = kj::getCaughtExceptionAsKj();
365    KJ_ASSERT(ex.getDescription().contains("closed stream"));
366  }
367 
368  // Should not throw, write should be a no-op
369}
370 
371KJ_TEST("newClosedWritableSink (end)") {
372  TestFixture fixture;
373  auto sink = newClosedWritableSink();
374 
375  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
376    // Test write on closed sink
377    co_await sink->end();
378  });
379 
380  // Should not throw, write should be a no-op
381}
382 
383KJ_TEST("newClosedWritableSink (aborted)") {
384  TestFixture fixture;
385  auto exception = KJ_EXCEPTION(FAILED, "test error");
386  auto sink = newClosedWritableSink();
387  sink->abort(exception.clone());
388 
389  try {
390    fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
391      // Test write on closed sink
392      co_await sink->end();
393    });
394    KJ_FAIL_REQUIRE("should have failed");
395  } catch (...) {
396    auto caught = kj::getCaughtExceptionAsKj();
397    KJ_ASSERT(caught.getDescription() == exception.getDescription());
398  }
399}
400 
401KJ_TEST("newErroredWritableSink (write)") {
402  TestFixture fixture;
403  auto exception = KJ_EXCEPTION(FAILED, "test error");
404  auto sink = newErroredWritableSink(exception.clone());
405 
406  try {
407    fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
408      // Test write on errored sink
409      co_await sink->write(kj::ArrayPtr<const kj::byte>());
410    });
411    KJ_FAIL_REQUIRE("should have failed");
412  } catch (...) {
413    auto caught = kj::getCaughtExceptionAsKj();
414    KJ_ASSERT(caught.getDescription() == exception.getDescription());
415  }
416}
417 
418KJ_TEST("newErroredWritableSink (end)") {
419  TestFixture fixture;
420  auto exception = KJ_EXCEPTION(FAILED, "test error");
421  auto sink = newErroredWritableSink(exception.clone());
422 
423  try {
424    fixture.runInIoContext(
425        [&](const auto& environment) -> kj::Promise<void> { co_await sink->end(); });
426    KJ_FAIL_REQUIRE("should have failed");
427  } catch (...) {
428    auto caught = kj::getCaughtExceptionAsKj();
429    KJ_ASSERT(caught.getDescription() == exception.getDescription());
430  }
431}
432 
433KJ_TEST("newNullWritableSink") {
434  TestFixture fixture;
435  auto sink = newNullWritableSink();
436 
437  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
438    // Test writing data
439    co_await sink->write(kj::ArrayPtr<const kj::byte>());
440    co_await sink->write(kj::ArrayPtr<const kj::byte>());
441    co_await sink->end();
442  });
443}
444 
445// ======================================================================================
446// Error Handling Tests
447 
448KJ_TEST("WritableSink error propagation") {
449  TestFixture fixture;
450  MemoryAsyncOutputStream inner;
451  inner.writeShouldError = true;
452  auto fakeOwn = kj::Own<MemoryAsyncOutputStream>(&inner, kj::NullDisposer::instance);
453  auto sink = newWritableSink(kj::mv(fakeOwn));
454  kj::byte testData[] = {1, 2, 3, 4, 5};
455 
456  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
457    try {
458      co_await sink->write(kj::ArrayPtr<const kj::byte>(testData, 5));
459      KJ_FAIL_REQUIRE("shuld have failed");
460    } catch (...) {
461      auto exception = kj::getCaughtExceptionAsKj();
462      KJ_ASSERT(exception.getDescription().contains("Expected failure"));
463    }
464    // The write should have been recorded.
465    KJ_ASSERT(inner.data.size() == 5);
466    // Let's make sure the error state is persistent.
467    try {
468      co_await sink->write(kj::ArrayPtr<const kj::byte>(testData, 5));
469      KJ_FAIL_REQUIRE("shuld have failed");
470    } catch (...) {
471      auto exception = kj::getCaughtExceptionAsKj();
472      KJ_ASSERT(exception.getDescription().contains("Expected failure"));
473    }
474    // No additional data should have been recorded.
475    KJ_ASSERT(inner.data.size() == 5);
476  });
477}
478 
479// ======================================================================================
480// Encoding Tests
481 
482KJ_TEST("WritableSink encoding responsibility transfer") {
483  TestFixture fixture;
484  MemoryAsyncOutputStream inner;
485  inner.writeShouldError = true;
486  auto fakeOwn = kj::Own<MemoryAsyncOutputStream>(&inner, kj::NullDisposer::instance);
487  auto sink = newWritableSink(kj::mv(fakeOwn));
488  KJ_ASSERT(sink->getEncoding() == rpc::StreamEncoding::IDENTITY);
489  KJ_ASSERT(sink->disownEncodingResponsibility() == rpc::StreamEncoding::IDENTITY);
490  KJ_ASSERT(sink->getEncoding() == rpc::StreamEncoding::IDENTITY);
491}
492 
493// ======================================================================================
494// Encoding-Aware WritableSink Implementation
495 
496KJ_TEST("Gzip-encoding sink") {
497  auto ctx = kj::setupAsyncIo();
498  TestFixture fixture({
499    .waitScope = ctx.waitScope,
500  });
501  MemoryAsyncOutputStream inner;
502  auto fakeOwn = kj::Own<MemoryAsyncOutputStream>(&inner, kj::NullDisposer::instance);
503  auto sink = newEncodedWritableSink(rpc::StreamEncoding::GZIP, kj::mv(fakeOwn));
504 
505  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
506    co_await sink->write("some data to gzip"_kjb);
507    co_await sink->end();
508 
509    auto mem = newMemoryInputStream(inner.data);
510    kj::GzipAsyncInputStream gunzip(*mem);
511    auto data = co_await gunzip.readAllText(kj::maxValue);
512    KJ_ASSERT(data == "some data to gzip"_kj);
513  });
514 
515  auto memInput = newMemoryInputStream(inner.data);
516  kj::GzipAsyncInputStream gunzip(*memInput);
517  auto decompressed = gunzip.readAllBytes().wait(ctx.waitScope);
518 
519  KJ_ASSERT(decompressed.asBytes() == "some data to gzip"_kjb);
520}
521 
522KJ_TEST("Gzip-encoding sink (identity)") {
523  TestFixture fixture;
524  MemoryAsyncOutputStream inner;
525  auto fakeOwn = kj::Own<MemoryAsyncOutputStream>(&inner, kj::NullDisposer::instance);
526  auto sink = newEncodedWritableSink(rpc::StreamEncoding::GZIP, kj::mv(fakeOwn));
527 
528  static const kj::byte check[] = {31, 139, 8, 0, 0, 0, 0, 0, 0, 3, 43, 206, 207, 77, 85, 72, 73,
529    44, 73, 84, 40, 201, 87, 72, 175, 202, 44, 0, 0, 40, 58, 113, 128, 17, 0, 0, 0};
530 
531  // When encoding is disowned, the data should be passed through unmodified.
532  sink->disownEncodingResponsibility();
533 
534  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
535    co_await sink->write(check);
536    co_await sink->end();
537  });
538 
539  KJ_ASSERT(inner.data == kj::arrayPtr(check, sizeof(check)));
540}
541 
542// ======================================================================================
543// IoContext-aware WritableSinkWrapper Tests
544 
545KJ_TEST("IoContext aware wrapper") {
546  TestFixture fixture;
547  MemoryAsyncOutputStream inner;
548  auto fakeOwn = kj::Own<MemoryAsyncOutputStream>(&inner, kj::NullDisposer::instance);
549  auto sink = newWritableSink(kj::mv(fakeOwn));
550 
551  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
552    IoContext& ioContext = environment.context;
553    auto wrapper = newIoContextWrappedWritableSink(ioContext, kj::mv(sink));
554    co_await wrapper->write("some data"_kjb);
555    co_await wrapper->end();
556  });
557 
558  KJ_ASSERT(inner.data == "some data"_kjb);
559}
560 
561// ======================================================================================
562// EndableAsyncOutputStream Tests
563 
564KJ_TEST("EndableAsyncOutputStream") {
565  TestFixture fixture;
566  MockEndable inner;
567  auto fakeOwn = kj::Own<MockEndable>(&inner, kj::NullDisposer::instance);
568  auto sink = newWritableSink(kj::mv(fakeOwn));
569 
570  fixture.runInIoContext([&](const auto& environment) -> kj::Promise<void> {
571    co_await sink->write("some data"_kjb);
572    co_await sink->end();
573  });
574 
575  KJ_ASSERT(inner.data == "some data"_kjb);
576  KJ_ASSERT(inner.isEnded);
577}
578 
579}  // namespace
580}  // namespace workerd::api::streams