1616using namespace ::testing;
1717using ::executorch::aten::ScalarType;
1818using ::executorch::aten::Tensor;
19+ using ::executorch::etdump::BufferDataSink;
1920using ::executorch::runtime::Error;
2021using ::executorch::runtime::Result;
2122using ::executorch::runtime::Span;
@@ -29,7 +30,8 @@ class BufferDataSinkTest : public ::testing::Test {
2930 buffer_size_ = 128 ; // Small size for testing
3031 buffer_ptr_ = malloc (buffer_size_);
3132 buffer_ = Span<uint8_t >(static_cast <uint8_t *>(buffer_ptr_), buffer_size_);
32- data_sink_ = std::make_unique<executorch::etdump::BufferDataSink>(buffer_);
33+ buffer_data_sink_ = std::make_unique<BufferDataSink>(
34+ std::move (BufferDataSink::create (buffer_).get ()));
3335 }
3436
3537 void TearDown () override {
@@ -39,11 +41,11 @@ class BufferDataSinkTest : public ::testing::Test {
3941 size_t buffer_size_;
4042 void * buffer_ptr_;
4143 Span<uint8_t > buffer_;
42- std::unique_ptr<executorch::etdump:: BufferDataSink> data_sink_ ;
44+ std::unique_ptr<BufferDataSink> buffer_data_sink_ ;
4345};
4446
4547TEST_F (BufferDataSinkTest, StorageSizeCheck) {
46- Result<size_t > ret = data_sink_ ->get_storage_size ();
48+ Result<size_t > ret = buffer_data_sink_ ->get_storage_size ();
4749 ASSERT_EQ (ret.error (), Error::Ok);
4850
4951 size_t storage_size = ret.get ();
@@ -55,7 +57,7 @@ TEST_F(BufferDataSinkTest, WriteOneTensorAndCheckData) {
5557 Tensor tensor = tf.make ({1 , 4 }, {1.0 , 2.0 , 3.0 , 4.0 });
5658
5759 Result<size_t > ret =
58- data_sink_ ->write (tensor.const_data_ptr (), tensor.nbytes ());
60+ buffer_data_sink_ ->write (tensor.const_data_ptr (), tensor.nbytes ());
5961 ASSERT_EQ (ret.error (), Error::Ok);
6062
6163 size_t offset = ret.get ();
@@ -75,9 +77,10 @@ TEST_F(BufferDataSinkTest, WriteMultiTensorsAndCheckData) {
7577 std::vector<Tensor> tensors = {
7678 tf.make ({1 , 4 }, {1.0 , 2.0 , 3.0 , 4.0 }),
7779 tf.make ({1 , 4 }, {5.0 , 6.0 , 7.0 , 8.0 })};
80+
7881 for (const auto & tensor : tensors) {
7982 Result<size_t > ret =
80- data_sink_ ->write (tensor.const_data_ptr (), tensor.nbytes ());
83+ buffer_data_sink_ ->write (tensor.const_data_ptr (), tensor.nbytes ());
8184 ASSERT_EQ (ret.error (), Error::Ok);
8285
8386 size_t offset = ret.get ();
@@ -94,8 +97,9 @@ TEST_F(BufferDataSinkTest, WriteMultiTensorsAndCheckData) {
9497TEST_F (BufferDataSinkTest, PointerAlignmentCheck) {
9598 TensorFactory<ScalarType::Float> tf;
9699 Tensor tensor = tf.make ({1 , 4 }, {1.0 , 2.0 , 3.0 , 4.0 });
100+
97101 Result<size_t > ret =
98- data_sink_ ->write (tensor.const_data_ptr (), tensor.nbytes ());
102+ buffer_data_sink_ ->write (tensor.const_data_ptr (), tensor.nbytes ());
99103 ASSERT_EQ (ret.error (), Error::Ok);
100104
101105 size_t offset = ret.get ();
@@ -112,12 +116,32 @@ TEST_F(BufferDataSinkTest, WriteUntilOverflow) {
112116 // Write tensors until we run out of space
113117 for (size_t i = 0 ; i < 2 ; i++) {
114118 Result<size_t > ret =
115- data_sink_ ->write (tensor.const_data_ptr (), tensor.nbytes ());
119+ buffer_data_sink_ ->write (tensor.const_data_ptr (), tensor.nbytes ());
116120 ASSERT_EQ (ret.error (), Error::Ok);
117121 }
118122
119123 // Attempting to write another tensor should raise an error
120124 Result<size_t > ret =
121- data_sink_ ->write (tensor.const_data_ptr (), tensor.nbytes ());
125+ buffer_data_sink_ ->write (tensor.const_data_ptr (), tensor.nbytes ());
122126 ASSERT_EQ (ret.error (), Error::OutOfResources);
123127}
128+
129+ TEST_F (BufferDataSinkTest, illegalAlignment) {
130+ // Create a buffer_data_sink_ with legal alignment that is a power of 2 and
131+ // greater than 0
132+ for (size_t i = 1 ; i <= 128 ; i <<= 1 ) {
133+ Result<BufferDataSink> buffer_data_sink_ret =
134+ BufferDataSink::create (buffer_, i);
135+ ASSERT_EQ (buffer_data_sink_ret.error (), Error::Ok);
136+ }
137+
138+ // Create a buffer_data_sink_ with illegal alignment that is not a power of 2
139+ // or greater than 0
140+ std::vector<size_t > illegal_alignments = {0 , 3 , 5 , 7 , 100 , 127 };
141+
142+ for (size_t i = 0 ; i < illegal_alignments.size (); i++) {
143+ Result<BufferDataSink> buffer_data_sink_ret =
144+ BufferDataSink::create (buffer_, illegal_alignments[i]);
145+ ASSERT_EQ (buffer_data_sink_ret.error (), Error::InvalidArgument);
146+ }
147+ }
0 commit comments