test_s3.py 1.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. import json
  2. import os
  3. import pytest
  4. from magic_pdf.data.data_reader_writer import S3DataReader, S3DataWriter
  5. @pytest.mark.skipif(
  6. os.getenv('S3_ACCESS_KEY', None) is None, reason='need s3 config!'
  7. )
  8. def test_multi_bucket_s3_reader_writer():
  9. """test multi bucket s3 reader writer must config s3 config in the
  10. environment export S3_BUCKET=xxx export S3_ACCESS_KEY=xxx export
  11. S3_SECRET_KEY=xxx export S3_ENDPOINT=xxx."""
  12. bucket = os.getenv('S3_BUCKET', '')
  13. ak = os.getenv('S3_ACCESS_KEY', '')
  14. sk = os.getenv('S3_SECRET_KEY', '')
  15. endpoint_url = os.getenv('S3_ENDPOINT', '')
  16. reader = S3DataReader(bucket=bucket, ak=ak, sk=sk, endpoint_url=endpoint_url)
  17. writer = S3DataWriter(bucket=bucket, ak=ak, sk=sk, endpoint_url=endpoint_url)
  18. bits = reader.read('meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl')
  19. assert bits == reader.read(
  20. f's3://{bucket}/meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl'
  21. )
  22. bits = reader.read(
  23. 'meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl?bytes=566,713'
  24. )
  25. assert bits == reader.read_at(
  26. 'meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl', 566, 713
  27. )
  28. assert len(json.loads(bits)) > 0
  29. writer.write_string(
  30. 'unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt', 'abc'
  31. )
  32. assert 'abc'.encode() == reader.read(
  33. 'unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt'
  34. )
  35. writer.write(
  36. f'{bucket}/unittest/data/data_reader_writer/multi_bucket_s3_data/test02.txt',
  37. '123'.encode(),
  38. )
  39. assert '123'.encode() == reader.read(
  40. 'unittest/data/data_reader_writer/multi_bucket_s3_data/test02.txt'
  41. )