import pytest
import os
from pyspark.context import SparkContext
from awsglue.context import GlueContext


@pytest.fixture(scope="session", autouse=True)
def scope_session():
    # テスト内で使い回せるようにS3のURLを環境変数に設定
    os.environ["TEST_S3_ENDPOINT_URL"] = "http://aws.local:4572"
    sc = SparkContext()
    # S3のエンドポイントをLocalStackへ差し替える
    sc._jsc.hadoopConfiguration().set("fs.s3a.endpoint", "http://aws.local:4572")
    sc._jsc.hadoopConfiguration().set("fs.s3a.path.style.access", "true")
    sc._jsc.hadoopConfiguration().set("fs.s3a.signing-algorithm", "S3SignerType")
    pytest.sc = sc
    pytest.glueContext = GlueContext(pytest.sc)
    pytest.spark = pytest.glueContext.spark_session

from awsglue.dynamicframe import DynamicFrame
from awsglue.context import GlueContext
import sys

def load_dynamic_frame_from_csv(glueContext: GlueContext, spark, bucket: str, path: str) -> DynamicFrame:
    p = "s3://{}/{}".format(bucket, path)
    return glueContext.create_dynamic_frame_from_options(
        connection_type="s3",
        connection_options={"paths": [p]},
        format="csv",
        format_options={"withHeader": True, "separator": ","},
    )

import pytest
import boto3
import json
import os
import sys
import io
import csv
from src.etl import load_dynamic_frame_from_csv
from botocore.client import Config

def test_load_dynamic_frame_from_csv():
    # setup
    inputs = [
      {
        "id": "1",
        "name": "xxx",
        "address": "xxx@example.co.jp"
      },
      {
        "id": "2",
        "name": "yyy",
        "address": "yyy@example.co.jp"
      }
    ]
    input_str = io.StringIO()
    w = csv.DictWriter(input_str, fieldnames=inputs[0].keys())
    w.writeheader()
    for input in inputs:
      w.writerow(input)
    s3 = boto3.resource(
        "s3",
        endpoint_url=os.environ["TEST_S3_ENDPOINT_URL"],
        region_name="ap-northeast-1",
        use_ssl=False,
        config=Config(s3={"addressing_style": "path"}),
    )
    bucket_name = "test-csv-bucket"
    bucket = s3.Bucket(bucket_name)
    bucket.create(ACL="public-read-write")
    body = input_str.getvalue()
    key = "user/2019/12/06/users.csv"
    bucket.put_object(Key=key, Body=body, ACL="public-read-write")

    # exec
    res_df = load_dynamic_frame_from_csv(pytest.glueContext, pytest.spark, bucket_name, key)

    # assert
    assert res_df.count() == len(inputs)
    res_df_json = res_df.toDF().toJSON().take(len(inputs))
    for res in res_df_json:
      r = json.loads(res)
      assert r in inputs