test_gcs.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. import numpy as np
  2. import pytest
  3. from pandas.compat import StringIO
  4. from pandas import DataFrame, date_range, read_csv
  5. from pandas.util import _test_decorators as td
  6. from pandas.util.testing import assert_frame_equal
  7. from pandas.io.common import is_gcs_url
  8. def test_is_gcs_url():
  9. assert is_gcs_url("gcs://pandas/somethingelse.com")
  10. assert is_gcs_url("gs://pandas/somethingelse.com")
  11. assert not is_gcs_url("s3://pandas/somethingelse.com")
  12. @td.skip_if_no('gcsfs')
  13. def test_read_csv_gcs(monkeypatch):
  14. df1 = DataFrame({'int': [1, 3], 'float': [2.0, np.nan], 'str': ['t', 's'],
  15. 'dt': date_range('2018-06-18', periods=2)})
  16. class MockGCSFileSystem():
  17. def open(*args):
  18. return StringIO(df1.to_csv(index=False))
  19. monkeypatch.setattr('gcsfs.GCSFileSystem', MockGCSFileSystem)
  20. df2 = read_csv('gs://test/test.csv', parse_dates=['dt'])
  21. assert_frame_equal(df1, df2)
  22. @td.skip_if_no('gcsfs')
  23. def test_to_csv_gcs(monkeypatch):
  24. df1 = DataFrame({'int': [1, 3], 'float': [2.0, np.nan], 'str': ['t', 's'],
  25. 'dt': date_range('2018-06-18', periods=2)})
  26. s = StringIO()
  27. class MockGCSFileSystem():
  28. def open(*args):
  29. return s
  30. monkeypatch.setattr('gcsfs.GCSFileSystem', MockGCSFileSystem)
  31. df1.to_csv('gs://test/test.csv', index=True)
  32. df2 = read_csv(StringIO(s.getvalue()), parse_dates=['dt'], index_col=0)
  33. assert_frame_equal(df1, df2)
  34. @td.skip_if_no('gcsfs')
  35. def test_gcs_get_filepath_or_buffer(monkeypatch):
  36. df1 = DataFrame({'int': [1, 3], 'float': [2.0, np.nan], 'str': ['t', 's'],
  37. 'dt': date_range('2018-06-18', periods=2)})
  38. def mock_get_filepath_or_buffer(*args, **kwargs):
  39. return (StringIO(df1.to_csv(index=False)),
  40. None, None, False)
  41. monkeypatch.setattr('pandas.io.gcs.get_filepath_or_buffer',
  42. mock_get_filepath_or_buffer)
  43. df2 = read_csv('gs://test/test.csv', parse_dates=['dt'])
  44. assert_frame_equal(df1, df2)
  45. @pytest.mark.skipif(td.safe_import('gcsfs'),
  46. reason='Only check when gcsfs not installed')
  47. def test_gcs_not_present_exception():
  48. with pytest.raises(ImportError) as e:
  49. read_csv('gs://test/test.csv')
  50. assert 'gcsfs library is required' in str(e.value)