From 72dbc8205bbc1711a39d127590dfa12b3d05c96a Mon Sep 17 00:00:00 2001 From: avogar Date: Mon, 9 Sep 2024 15:12:17 +0000 Subject: [PATCH] Fix unit tests --- src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp | 9 +++++++++ src/IO/WriteBufferFromS3.cpp | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp b/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp index cbcfe0bdb97..229b36a05f6 100644 --- a/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp +++ b/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp @@ -251,7 +251,16 @@ void WriteBufferFromAzureBlobStorage::allocateBuffer() /// First buffer was already allocated in BufferWithOwnMemory constructor with buffer size provided in constructor. /// It will be reallocated in subsequent nextImpl calls up to the desired buffer size from buffer_allocation_policy. if (buffer_allocation_policy->getBufferNumber() == 1) + { + /// Reduce memory size if initial size was larger then desired size from buffer_allocation_policy. + /// Usually it doesn't happen but we have it in unit tests. + if (memory.size() > buffer_allocation_policy->getBufferSize()) + { + memory.resize(buffer_allocation_policy->getBufferSize()); + WriteBuffer::set(memory.data(), memory.size()); + } return; + } auto size = buffer_allocation_policy->getBufferSize(); memory = Memory(size); diff --git a/src/IO/WriteBufferFromS3.cpp b/src/IO/WriteBufferFromS3.cpp index 424708cf2c8..3a4ab3ee882 100644 --- a/src/IO/WriteBufferFromS3.cpp +++ b/src/IO/WriteBufferFromS3.cpp @@ -355,7 +355,16 @@ void WriteBufferFromS3::allocateBuffer() /// First buffer was already allocated in BufferWithOwnMemory constructor with provided in constructor buffer size. /// It will be reallocated in subsequent nextImpl calls up to the desired buffer size from buffer_allocation_policy. if (buffer_allocation_policy->getBufferNumber() == 1) + { + /// Reduce memory size if initial size was larger then desired size from buffer_allocation_policy. + /// Usually it doesn't happen but we have it in unit tests. + if (memory.size() > buffer_allocation_policy->getBufferSize()) + { + memory.resize(buffer_allocation_policy->getBufferSize()); + WriteBuffer::set(memory.data(), memory.size()); + } return; + } memory = Memory(buffer_allocation_policy->getBufferSize()); WriteBuffer::set(memory.data(), memory.size());