+#include <stdio.h>
+#include <string.h>
+#include <assert.h>
+
+#include "bytesource.h"
+#include "choice.h"
+
+struct custom_read_userdata {
+ uint8_t* bytes;
+ unsigned bytes_left;
+};
+
+ssize_t custom_read(void* userdata, uint8_t* buf, size_t count)
+{
+ struct custom_read_userdata* ud = (struct custom_read_userdata*)userdata;
+ size_t num_to_read = (ud->bytes_left > count ? count : ud->bytes_left);
+ memcpy(buf, ud->bytes, num_to_read);
+ ud->bytes += num_to_read;
+ ud->bytes_left -= num_to_read;
+ return num_to_read;
+}
+
+ssize_t custom_read_slow(void* userdata, uint8_t* buf, size_t count)
+{
+ struct custom_read_userdata* ud = (struct custom_read_userdata*)userdata;
+ size_t num_to_read = ((count > 0 && ud->bytes_left > 0) ? 1 : 0);
+ memcpy(buf, ud->bytes, num_to_read);
+ ud->bytes += num_to_read;
+ ud->bytes_left -= num_to_read;
+ return num_to_read;
+}
+
+// Two streams, separated by a marker.
+uint8_t bytes[] = {
+ 0x01, 0x02, 0x03, 0xff, 0x00, 0x04, // some bytes
+ 0xff, 0xf7, // a marker
+ 0x05, 0x06, 0x07, 0x08, 0xff, 0x00 // more bytes
+};
+
+// Expected data.
+uint8_t first_stream[] = {
+ 0x01, 0x02, 0x03, 0xff, 0x04
+};
+uint8_t second_stream[] = {
+ 0x05, 0x06, 0x07, 0x08, 0xff
+};
+
+// Reading with a regular source.
+void test_basic_reading()
+{
+ struct custom_read_userdata ud;
+ ud.bytes = bytes;
+ ud.bytes_left = sizeof(bytes);
+
+ struct byte_source source;
+ init_byte_source(&source, custom_read, &ud);
+
+ // Read the first stream.
+ uint8_t buf[4096];
+ ssize_t ret;
+ ret = byte_source_input_func(&source, buf, 4096);
+ assert(ret == sizeof(first_stream));
+ assert(memcmp(buf, first_stream, sizeof(first_stream)) == 0);
+
+ // Now we should get EOF.
+ ret = byte_source_input_func(&source, buf, 4096);
+ assert(ret == 0);
+
+ // Read the marker.
+ uint8_t marker = byte_source_read_marker(&source);
+ assert(marker == 0xf7);
+
+ // Read the second stream.
+ ret = byte_source_input_func(&source, buf, 4096);
+ assert(ret == sizeof(second_stream));
+ assert(memcmp(buf, second_stream, sizeof(second_stream)) == 0);
+
+ // ...and EOF again.
+ ret = byte_source_input_func(&source, buf, 4096);
+ assert(ret == 0);
+}
+
+// Reading with a slow source.
+void test_slow_source()
+{
+ struct custom_read_userdata ud;
+ ud.bytes = bytes;
+ ud.bytes_left = sizeof(bytes);
+
+ struct byte_source source;
+ init_byte_source(&source, custom_read_slow, &ud);
+
+ uint8_t buf[4096];
+ ssize_t ret;
+
+ // Read the first stream. Since the source is slow, we'll get one by
+ // one byte, even though we asked for 4096.
+ for (unsigned i = 0; i < sizeof(first_stream); ++i) {
+ ret = byte_source_input_func(&source, buf, 4096);
+ assert(ret == 1);
+ assert(buf[0] == first_stream[i]);
+ }
+
+ // Now we should get EOF.
+ ret = byte_source_input_func(&source, buf, 4096);
+ assert(ret == 0);
+
+ // Read the marker.
+ uint8_t marker = byte_source_read_marker(&source);
+ assert(marker == 0xf7);
+
+ // Read the second stream.
+ for (unsigned i = 0; i < sizeof(second_stream); ++i) {
+ uint8_t buf[4096];
+ ssize_t ret;
+ ret = byte_source_input_func(&source, buf, 4096);
+ assert(ret == 1);
+ assert(buf[0] == second_stream[i]);
+ }
+
+ // ...and EOF again.
+ ret = byte_source_input_func(&source, buf, 4096);
+ assert(ret == 0);
+}
+
+int main(void)
+{
+ init_choices();
+
+ printf("test_basic_reading()\n");
+ test_basic_reading();
+
+ printf("test_slow_source()\n");
+ test_slow_source();
+
+ printf("All tests pass.\n");
+ return 0;
+}