Add tests for Buffer and URL I/O

Author: TomoyukiSuzuki-Aidea

pandas/tests/io/test_pickle.py

@@ -22,6 +22,7 @@
 import pytest
 
 from pandas.compat import _get_lzma_file, _import_lzma, is_platform_little_endian
+import pandas.util._test_decorators as td
 
 import pandas as pd
 from pandas import Index
@@ -390,3 +391,92 @@ def test_unicode_decode_error():
     # just test the columns are correct since the values are random
     excols = pd.Index(["a", "b", "c"])
     tm.assert_index_equal(df.columns, excols)
+
+
+# ---------------------
+# tests for buffer I/O
+# ---------------------
+
+
+def test_pickle_buffer_roundtrip():
+    with tm.ensure_clean() as path:
+        df = tm.makeDataFrame()
+        df.to_pickle(open(path, "wb"))
+        result = pd.read_pickle(open(path, "rb"))
+        tm.assert_frame_equal(df, result)
+
+
+# ---------------------
+# tests for URL I/O
+# ---------------------
+
+
+@pytest.mark.parametrize("mockurl", ["http://url.com", "ftp://test.com"])
+def test_pickle_generalurl_read(monkeypatch, mockurl):
+    def python_pickler(obj, path):
+        with open(path, "wb") as fh:
+            pickle.dump(obj, fh, protocol=-1)
+
+    class MockReadResponse:
+        def __init__(self, path):
+            self.file = open(path, "rb")
+            self.headers = {"Content-Encoding": None}
+
+        def read(self):
+            return self.file.read()
+
+        def close(self):
+            return self.file.close()
+
+    with tm.ensure_clean() as path:
+
+        def mock_urlopen_read(*args, **kwargs):
+            return MockReadResponse(path)
+
+        df = tm.makeDataFrame()
+        python_pickler(df, path)
+        monkeypatch.setattr("urllib.request.urlopen", mock_urlopen_read)
+        result = pd.read_pickle(mockurl)
+        tm.assert_frame_equal(df, result)
+
+
+@td.skip_if_no("gcsfs")
+@pytest.mark.parametrize("mockurl", ["gs://gcs.com", "gcs://gcs.com"])
+def test_pickle_gcsurl_roundtrip(monkeypatch, mockurl):
+    with tm.ensure_clean() as path:
+
+        class MockGCSFileSystem:
+            def __init__(self, *args, **kwargs):
+                pass
+
+            def open(self, *args):
+                mode = args[1] or None
+                f = open(path, mode)
+                return f
+
+        monkeypatch.setattr("gcsfs.GCSFileSystem", MockGCSFileSystem)
+        df = tm.makeDataFrame()
+        df.to_pickle(mockurl)
+        result = pd.read_pickle(mockurl)
+        tm.assert_frame_equal(df, result)
+
+
+@td.skip_if_no("s3fs")
+@pytest.mark.parametrize("mockurl", ["s3://s3.com", "s3n://s3.com", "s3a://s3.com"])
+def test_pickle_s3url_roundtrip(monkeypatch, mockurl):
+    with tm.ensure_clean() as path:
+
+        class MockS3FileSystem:
+            def __init__(self, *args, **kwargs):
+                pass
+
+            def open(self, *args):
+                mode = args[1] or None
+                f = open(path, mode)
+                return f
+
+        monkeypatch.setattr("s3fs.S3FileSystem", MockS3FileSystem)
+        df = tm.makeDataFrame()
+        df.to_pickle(mockurl)
+        result = pd.read_pickle(mockurl)
+        tm.assert_frame_equal(df, result)