github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/io/filesystemio_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Tests for filesystemio."""

# pytype: skip-file

import io
import logging
import multiprocessing
import os
import threading
import unittest

from apache_beam.io import filesystemio

_LOGGER = logging.getLogger(__name__)


class FakeDownloader(filesystemio.Downloader):
  def __init__(self, data):
    self._data = data
    self.last_read_size = -1

  @property
  def size(self):
    return len(self._data)

  def get_range(self, start, end):
    self.last_read_size = end - start
    return self._data[start:end]


class FakeUploader(filesystemio.Uploader):
  def __init__(self):
    self.data = b''
    self.last_write_size = -1
    self.finished = False

  def last_error(self):
    return None

  def put(self, data):
    assert not self.finished
    self.data += data.tobytes()
    self.last_write_size = len(data)

  def finish(self):
    self.finished = True


class TestDownloaderStream(unittest.TestCase):
  def test_file_attributes(self):
    downloader = FakeDownloader(data=None)
    stream = filesystemio.DownloaderStream(downloader)
    self.assertEqual(stream.mode, 'rb')
    self.assertTrue(stream.readable())
    self.assertFalse(stream.writable())
    self.assertTrue(stream.seekable())

  def test_read_empty(self):
    downloader = FakeDownloader(data=b'')
    stream = filesystemio.DownloaderStream(downloader)
    self.assertEqual(stream.read(), b'')

  def test_read(self):
    data = b'abcde'
    downloader = FakeDownloader(data)
    stream = filesystemio.DownloaderStream(downloader)

    # Read size is exactly what was passed to read() (unbuffered).
    self.assertEqual(stream.read(1), data[0:1])
    self.assertEqual(downloader.last_read_size, 1)
    self.assertEqual(stream.read(), data[1:])
    self.assertEqual(downloader.last_read_size, len(data) - 1)

  def test_read_buffered(self):
    data = b'abcde'
    downloader = FakeDownloader(data)
    buffer_size = 2
    stream = io.BufferedReader(
        filesystemio.DownloaderStream(downloader), buffer_size)

    # Verify that buffering works and is reading ahead.
    self.assertEqual(stream.read(1), data[0:1])
    self.assertEqual(downloader.last_read_size, buffer_size)
    self.assertEqual(stream.read(), data[1:])


class TestUploaderStream(unittest.TestCase):
  def test_file_attributes(self):
    uploader = FakeUploader()
    stream = filesystemio.UploaderStream(uploader)
    self.assertEqual(stream.mode, 'wb')
    self.assertFalse(stream.readable())
    self.assertTrue(stream.writable())
    self.assertFalse(stream.seekable())

  def test_write_empty(self):
    uploader = FakeUploader()
    stream = filesystemio.UploaderStream(uploader)
    data = b''
    stream.write(memoryview(data))
    self.assertEqual(uploader.data, data)

  def test_write(self):
    data = b'abcde'
    uploader = FakeUploader()
    stream = filesystemio.UploaderStream(uploader)

    # Unbuffered writes.
    stream.write(memoryview(data[0:1]))
    self.assertEqual(uploader.data[0], data[0])
    self.assertEqual(uploader.last_write_size, 1)
    stream.write(memoryview(data[1:]))
    self.assertEqual(uploader.data, data)
    self.assertEqual(uploader.last_write_size, len(data) - 1)

  def test_write_buffered(self):
    data = b'abcde'
    uploader = FakeUploader()
    buffer_size = 2
    stream = io.BufferedWriter(
        filesystemio.UploaderStream(uploader), buffer_size)

    # Verify that buffering works: doesn't write to uploader until buffer is
    # filled.
    stream.write(data[0:1])
    self.assertEqual(-1, uploader.last_write_size)
    stream.write(data[1:])
    stream.close()
    self.assertEqual(data, uploader.data)


class TestPipeStream(unittest.TestCase):
  def _read_and_verify(self, stream, expected, buffer_size, success):
    data_list = []
    bytes_read = 0
    seen_last_block = False
    while True:
      data = stream.read(buffer_size)
      self.assertLessEqual(len(data), buffer_size)
      if len(data) < buffer_size:
        # Test the constraint that the pipe stream returns less than the buffer
        # size only when at the end of the stream.
        if data:
          self.assertFalse(seen_last_block)
        seen_last_block = True
      if not data:
        break
      data_list.append(data)
      bytes_read += len(data)
      self.assertEqual(stream.tell(), bytes_read)
    self.assertEqual(b''.join(data_list), expected)
    success[0] = True

  def _read_and_seek(self, stream, expected, buffer_size, success):
    data_list = []
    bytes_read = 0
    while True:
      data = stream.read(buffer_size)

      # Test bad seek positions.
      with self.assertRaises(NotImplementedError):
        stream.seek(bytes_read + 1)
      with self.assertRaises(NotImplementedError):
        stream.seek(bytes_read - 1)

      # Rewind stream and test that it reads back the same data again.
      stream.seek(bytes_read)
      data2 = stream.read(buffer_size)
      self.assertEqual(data, data2)

      if not data:
        break
      data_list.append(data)
      bytes_read += len(data)
      self.assertEqual(stream.tell(), bytes_read)
    self.assertEqual(len(b''.join(data_list)), len(expected))
    self.assertEqual(b''.join(data_list), expected)
    success[0] = True

  def test_pipe_stream(self):
    block_sizes = list(4**i for i in range(0, 12))
    data_blocks = list(os.urandom(size) for size in block_sizes)
    expected = b''.join(data_blocks)

    buffer_sizes = [100001, 512 * 1024, 1024 * 1024]

    for buffer_size in buffer_sizes:
      for target in [self._read_and_verify, self._read_and_seek]:
        _LOGGER.info('buffer_size=%s, target=%s' % (buffer_size, target))
        parent_conn, child_conn = multiprocessing.Pipe()
        stream = filesystemio.PipeStream(child_conn)
        success = [False]
        child_thread = threading.Thread(
            target=target, args=(stream, expected, buffer_size, success))
        child_thread.start()
        for data in data_blocks:
          parent_conn.send_bytes(data)
        parent_conn.close()
        child_thread.join()
        self.assertTrue(success[0], 'error in test thread')

  def test_pipe_stream_rewind_buffer(self):
    buffer_size = 512
    data = os.urandom(buffer_size)

    parent_conn, child_conn = multiprocessing.Pipe()
    parent_conn.send_bytes(data)
    parent_conn.close()
    stream = filesystemio.PipeStream(child_conn)

    # Regular read.
    read_data = stream.read(buffer_size)
    self.assertEqual(data, read_data)

    # Rewind buffer_size bytes.
    stream.seek(0)
    read_data = stream.read(buffer_size)
    self.assertEqual(data, read_data)

    # Read 0 bytes. Rewind buffer still points to offset 0.
    read_data = stream.read(buffer_size)
    self.assertFalse(read_data)
    stream.seek(0)
    read_data = stream.read(buffer_size)
    self.assertEqual(data, read_data)


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()