From 54b33862449d265ea7539d453f2c238a91f5cefc Mon Sep 17 00:00:00 2001
From: igorborgest <igorborgest@gmail.com>
Date: Tue, 19 May 2020 16:17:32 -0300
Subject: [PATCH] Add tests to the new default session. #245

---
 awswrangler/_utils.py                    |  4 +-
 testing/test_awswrangler/test_session.py | 20 +++++++
 tutorials/02 - Sessions.ipynb            | 70 +++++++++++++++++++++---
 3 files changed, 85 insertions(+), 9 deletions(-)
 create mode 100644 testing/test_awswrangler/test_session.py

diff --git a/awswrangler/_utils.py b/awswrangler/_utils.py
index bcb8eab79..b31efd803 100644
--- a/awswrangler/_utils.py
+++ b/awswrangler/_utils.py
@@ -23,7 +23,9 @@ def ensure_session(session: Optional[boto3.Session] = None) -> boto3.Session:
         return session
     # Ensure the boto3's default session is used so that its parameters can be
     # set via boto3.setup_default_session()
-    return boto3._get_default_session()  # pylint: disable=protected-access
+    if boto3.DEFAULT_SESSION is not None:
+        return boto3.DEFAULT_SESSION
+    return boto3.Session()
 
 
 def client(service_name: str, session: Optional[boto3.Session] = None) -> boto3.client:
diff --git a/testing/test_awswrangler/test_session.py b/testing/test_awswrangler/test_session.py
new file mode 100644
index 000000000..172f97bba
--- /dev/null
+++ b/testing/test_awswrangler/test_session.py
@@ -0,0 +1,20 @@
+import logging
+
+import boto3
+
+import awswrangler as wr
+
+logging.basicConfig(level=logging.INFO, format="[%(asctime)s][%(levelname)s][%(name)s][%(funcName)s] %(message)s")
+logging.getLogger("awswrangler").setLevel(logging.DEBUG)
+logging.getLogger("botocore.credentials").setLevel(logging.CRITICAL)
+
+
+def test_default_session():
+    boto3.setup_default_session(region_name="us-east-1")
+    assert wr._utils.ensure_session().region_name == "us-east-1"
+    boto3.setup_default_session(region_name="us-east-2")
+    assert wr._utils.ensure_session().region_name == "us-east-2"
+    boto3.setup_default_session(region_name="us-west-1")
+    assert wr._utils.ensure_session().region_name == "us-west-1"
+    boto3.setup_default_session(region_name="us-west-2")
+    assert wr._utils.ensure_session().region_name == "us-west-2"
diff --git a/tutorials/02 - Sessions.ipynb b/tutorials/02 - Sessions.ipynb
index 97b376968..07ace8f94 100644
--- a/tutorials/02 - Sessions.ipynb	
+++ b/tutorials/02 - Sessions.ipynb	
@@ -19,7 +19,7 @@
     "\n",
     "Wrangler will not store any kind of state internally, and users is in charge of all the Sessions management, if necessary.\n",
     "\n",
-    "Most Wrangler functions receive the optional `boto3_session` argument. If None is received, a default boto3 Session will be temporary created to run the function."
+    "Most Wrangler functions receive the optional `boto3_session` argument. If None is received, the default boto3 Session will be used."
    ]
   },
   {
@@ -36,7 +36,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Using the default Sessions"
+    "## Using the default Session"
    ]
   },
   {
@@ -63,7 +63,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Using custom Sessions"
+    "## Customizing and using the default Session"
    ]
   },
   {
@@ -83,7 +83,61 @@
     }
    ],
    "source": [
-    "wr.s3.does_object_exist(\"s3://noaa-ghcn-pds/fake\", boto3_session=boto3.Session(region_name=\"us-east-2\"))"
+    "boto3.setup_default_session(region_name=\"us-east-2\")\n",
+    "\n",
+    "wr.s3.does_object_exist(\"s3://noaa-ghcn-pds/fake\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Using a new custom Session"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "my_session = boto3.Session(region_name=\"us-east-2\")\n",
+    "\n",
+    "wr.s3.does_object_exist(\"s3://noaa-ghcn-pds/fake\", boto3_session=my_session)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "False"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "my_session = boto3.Session(region_name=\"us-east-2\")\n",
+    "\n",
+    "wr.s3.does_object_exist(\"s3://noaa-ghcn-pds/fake\", boto3_session=my_session)"
    ]
   },
   {
@@ -96,9 +150,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "conda_python3",
+   "display_name": "Python 3",
    "language": "python",
-   "name": "conda_python3"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -110,9 +164,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.5"
+   "version": "3.6.10"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
\ No newline at end of file