(apache#25316) update tests (draft)

wrossmorrow · Feb 11, 2023 · 641d3f3 · 641d3f3
1 parent 4991413
commit 641d3f3
Showing 1 changed file with 19 additions and 6 deletions.
diff --git a/sdks/python/apache_beam/io/filesystem_test.py b/sdks/python/apache_beam/io/filesystem_test.py
@@ -29,6 +29,7 @@
 import tempfile
 import unittest
 import zlib
+import lzma
 from io import BytesIO
 
 import zstandard
@@ -316,6 +317,10 @@ def _create_compressed_file(self, compression_type, content):
       compress_open = zstandard.open
       with compress_open(file_name, 'wb') as f:
         f.write(content)
+    elif compression_type == CompressionTypes.LZMA:
+      compress_open = lzma.open
+      with compress_open(file_name, 'wb') as f:
+        f.write(content)
     else:
       assert False, "Invalid compression type: %s" % compression_type
 
@@ -340,7 +345,8 @@ def test_seek_set(self):
     for compression_type in [CompressionTypes.BZIP2,
                              CompressionTypes.DEFLATE,
                              CompressionTypes.GZIP,
-                             CompressionTypes.ZSTD]:
+                             CompressionTypes.ZSTD,
+                             CompressionTypes.LZMA]:
       file_name = self._create_compressed_file(compression_type, self.content)
       with open(file_name, 'rb') as f:
         compressed_fd = CompressedFile(
@@ -375,7 +381,8 @@ def test_seek_cur(self):
     for compression_type in [CompressionTypes.BZIP2,
                              CompressionTypes.DEFLATE,
                              CompressionTypes.GZIP,
-                             CompressionTypes.ZSTD]:
+                             CompressionTypes.ZSTD,
+                             CompressionTypes.LZMA]:
       file_name = self._create_compressed_file(compression_type, self.content)
       with open(file_name, 'rb') as f:
         compressed_fd = CompressedFile(
@@ -410,7 +417,8 @@ def test_read_from_end_returns_no_data(self):
     for compression_type in [CompressionTypes.BZIP2,
                              CompressionTypes.DEFLATE,
                              CompressionTypes.GZIP,
-                             CompressionTypes.ZSTD]:
+                             CompressionTypes.ZSTD,
+                             CompressionTypes.LZMA]:
       file_name = self._create_compressed_file(compression_type, self.content)
       with open(file_name, 'rb') as f:
         compressed_fd = CompressedFile(
@@ -428,7 +436,8 @@ def test_seek_outside(self):
     for compression_type in [CompressionTypes.BZIP2,
                              CompressionTypes.DEFLATE,
                              CompressionTypes.GZIP,
-                             CompressionTypes.ZSTD]:
+                             CompressionTypes.ZSTD,
+                             CompressionTypes.LZMA]:
       file_name = self._create_compressed_file(compression_type, self.content)
       with open(file_name, 'rb') as f:
         compressed_fd = CompressedFile(
@@ -453,7 +462,8 @@ def test_read_and_seek_back_to_beginning(self):
     for compression_type in [CompressionTypes.BZIP2,
                              CompressionTypes.DEFLATE,
                              CompressionTypes.GZIP,
-                             CompressionTypes.ZSTD]:
+                             CompressionTypes.ZSTD,
+                             CompressionTypes.LZMA]:
       file_name = self._create_compressed_file(compression_type, self.content)
       with open(file_name, 'rb') as f:
         compressed_fd = CompressedFile(
@@ -520,6 +530,8 @@ def create_test_file(compression_type, lines):
         compress_factory = gzip.open
       elif compression_type == CompressionTypes.ZSTD:
         compress_factory = zstandard.open
+      elif compression_type == CompressionTypes.LZMA:
+        compress_factory = lzma.open
       else:
         assert False, "Invalid compression type: %s" % compression_type
       for line in lines:
@@ -547,7 +559,8 @@ def timeout_handler():
       test_lines = tuple(generate_random_line() for i in range(num_test_lines))
       for compression_type in [CompressionTypes.BZIP2,
                                CompressionTypes.GZIP,
-                               CompressionTypes.ZSTD]:
+                               CompressionTypes.ZSTD,
+                               CompressionTypes.LZMA]:
         file_name = create_test_file(compression_type, test_lines)
         timer.start()
         with open(file_name, 'rb') as f: