text: # Copyright (c) 2020, NVIDIA CORPORATION.
  
  import os
  import shlex
  import subprocess
  import time
  from contextlib import contextmanager
  from io import BytesIO
  
  import numpy as np
  import pandas as pd
  import pyarrow as pa
  import pyarrow.orc
  import pytest
  
  import cudf
  from cudf.tests.utils import assert_eq
  
  moto = pytest.importorskip("moto", minversion="1.3.14")
  boto3 = pytest.importorskip("boto3")
  requests = pytest.importorskip("requests")
  s3fs = pytest.importorskip("s3fs")
  
  
  @contextmanager
  def ensure_safe_environment_variables():
      """
      Get a context manager to safely set environment variables
      All changes will be undone on close, hence environment variables set
      within this contextmanager will neither persist nor change global state.
      """
      saved_environ = dict(os.environ)
      try:
          yield
      finally:
          os.environ.clear()
          os.environ.update(saved_environ)
  
  
  @pytest.fixture(scope="session")
  def s3_base(worker_id):
      """
      Fixture to set up moto server in separate process
      """
      with ensure_safe_environment_variables():
          # Fake aws credentials exported to prevent botocore looking for
          # system aws credentials, https://github.com/spulec/moto/issues/1793
          os.environ.setdefault("AWS_ACCESS_KEY_ID", "foobar_key")
          os.environ.setdefault("AWS_SECRET_ACCESS_KEY", "foobar_secret")
  
          # Launching moto in server mode, i.e., as a separate process
          # with an S3 endpoint on localhost
  
          endpoint_port = (
              5000
              if worker_id == "master"
              else 5550 + int(worker_id.lstrip("gw"))
          )
          endpoint_uri = f"http://127.0.0.1:{endpoint_port}/"
  
          proc = subprocess.Popen(
              shlex.split(f"moto_server s3 -p {endpoint_port}"),
          )
  
          timeout = 5
          while timeout > 0:
              try:
                  # OK to go once server is accepting connections
                  r = requests.get(endpoint_uri)
                  if r.ok:
                      break
              except Exception:
                  pass
              timeout -= 0.1
              time.sleep(0.1)
          yield endpoint_uri
  
          proc.terminate()
          proc.wait()
  
  
  @pytest.fixture()
  def s3so(worker_id):
      """
      Returns s3 storage options to pass to fsspec
      """
      endpoint_port = (
          5000 if worker_id == "master" else 5550 + int(worker_id.lstrip("gw"))
      )
      endpoint_uri = f"http://127.0.0.1:{endpoint_port}/"
  
      return {"client_kwargs": {"endpoint_url": endpoint_uri}}
  
  
  @contextmanager
  def s3_context(s3_base, bucket, files=None):
      if files is None:
          files = {}
      with ensure_safe_environment_variables():
          client = boto3.client("s3", endpoint_url=s3_base)
          client.create_bucket(Bucket=bucket, ACL="public-read-write")
          for f, data in files.items():
              client.put_object(Bucket=bucket, Key=f, Body=data)
  
          yield s3fs.S3FileSystem(client_kwargs={"endpoint_url": s3_base})
  
          for f, data in files.items():
              try:
                  client.delete_object(Bucket=bucket, Key=f)
              except Exception:
                  pass
  
  
  @pytest.fixture
  def pdf(scope="module"):
      df = pd.DataFrame()
      df["Integer"] = np.array([2345, 11987, 9027, 9027])
      df["Float"] = np.array([9.001, 8.343, 6, 2.781])
      df["Integer2"] = np.array([2345, 106, 2088, 789277])
      df["String"] = np.array(["Alpha", "Beta", "Gamma", "Delta"])
      df["Boolean"] = np.array([True, False, True, False])
      return df
  
  
  def test_read_csv(s3_base, s3so, pdf):
      # Write to buffer
      fname = "test_csv_reader.csv"
      bname = "csv"
      buffer = pdf.to_csv(index=False)
      with s3_context(s3_base=s3_base, bucket=bname, files={fname: buffer}):
          got = cudf.read_csv(
              "s3://{}/{}".format(bname, fname), storage_options=s3so
          )
  
      assert_eq(pdf, got)
  
  
  @pytest.mark.parametrize("chunksize", [None, 3])
  def test_write_csv(s3_base, s3so, pdf, chunksize):
      # Write to buffer
      fname = "test_csv_writer.csv"
      bname = "csv"
      gdf = cudf.from_pandas(pdf)
      with s3_context(s3_base=s3_base, bucket=bname) as s3fs:
          gdf.to_csv(
              "s3://{}/{}".format(bname, fname),
              index=False,
              chunksize=chunksize,
              storage_options=s3so,
          )
          assert s3fs.exists("s3://{}/{}".format(bname, fname))
  
          # TODO: Update to use `storage_options` from pandas v1.2.0
          got = pd.read_csv(s3fs.open("s3://{}/{}".format(bname, fname)))
  
      assert_eq(pdf, got)
  
  
  def test_read_parquet(s3_base, s3so, pdf):
      fname = "test_parquet_reader.parquet"
      bname = "parquet"
      buffer = BytesIO()
      pdf.to_parquet(path=buffer)
      buffer.seek(0)
      with s3_context(s3_base=s3_base, bucket=bname, files={fname: buffer}):
          got = cudf.read_parquet(
              "s3://{}/{}".format(bname, fname), storage_options=s3so
          )
  
      assert_eq(pdf, got)
  
  
  def test_write_parquet(s3_base, s3so, pdf):
      fname = "test_parquet_writer.parquet"
      bname = "parquet"
      gdf = cudf.from_pandas(pdf)
      with s3_context(s3_base=s3_base, bucket=bname) as s3fs:
          gdf.to_parquet("s3://{}/{}".format(bname, fname), storage_options=s3so)
          assert s3fs.exists("s3://{}/{}".format(bname, fname))
  
          got = pd.read_parquet(s3fs.open("s3://{}/{}".format(bname, fname)))
  
      assert_eq(pdf, got)
  
  
  def test_read_json(s3_base, s3so):
      fname = "test_json_reader.json"
      bname = "json"
      buffer = (
          b'{"amount": 100, "name": "Alice"}\n'
          b'{"amount": 200, "name": "Bob"}\n'
          b'{"amount": 300, "name": "Charlie"}\n'
          b'{"amount": 400, "name": "Dennis"}\n'
      )
  
      with s3_context(s3_base=s3_base, bucket=bname, files={fname: buffer}):
          got = cudf.read_json(
              "s3://{}/{}".format(bname, fname),
              engine="cudf",
              orient="records",
              lines=True,
              storage_options=s3so,
          )
  
      expect = pd.read_json(buffer, lines=True)
      assert_eq(expect, got)
  
  
  def test_read_orc(s3_base, s3so, datadir):
      source_file = str(datadir / "orc" / "TestOrcFile.testSnappy.orc")
      fname = "test_orc_reader.orc"
      bname = "orc"
      expect = pa.orc.ORCFile(source_file).read().to_pandas()
  
      with open(source_file, "rb") as f:
          buffer = f.read()
  
      with s3_context(s3_base=s3_base, bucket=bname, files={fname: buffer}):
          got = cudf.read_orc(
              "s3://{}/{}".format(bname, fname), storage_options=s3so
          )
  
      assert_eq(expect, got)
  
  
  def test_write_orc(s3_base, s3so, pdf):
      fname = "test_orc_writer.orc"
      bname = "orc"
      gdf = cudf.from_pandas(pdf)
      with s3_context(s3_base=s3_base, bucket=bname) as s3fs:
          gdf.to_orc("s3://{}/{}".format(bname, fname), storage_options=s3so)
          assert s3fs.exists("s3://{}/{}".format(bname, fname))
  
          with s3fs.open("s3://{}/{}".format(bname, fname)) as f:
              got = pa.orc.ORCFile(f).read().to_pandas()
  
      assert_eq(pdf, got)
  