Added tests for S3Client with provided session

test/contrib/s3_test.py

@@ -41,6 +41,7 @@
 
 AWS_ACCESS_KEY = "XXXXXXXXXXXXXXXXXXXX"
 AWS_SECRET_KEY = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
+AWS_SESSION_TOKEN = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
 
 
 def create_bucket():
@@ -49,7 +50,6 @@ def create_bucket():
     conn.create_bucket(Bucket='mybucket')
     return conn
 
-
 @attr('aws')
 class TestS3Target(unittest.TestCase, FileSystemTargetTestMixin):
 
@@ -72,6 +72,11 @@ def create_target(self, format=None, **kwargs):
         create_bucket()
         return S3Target('s3://mybucket/test_file', client=client, format=format, **kwargs)
 
+    def create_target_with_session(self, format=None, **kwargs):
+        client = S3Client(AWS_ACCESS_KEY, AWS_SECRET_KEY, AWS_SESSION_TOKEN)
+        create_bucket()
+        return S3Target('s3://mybucket/test_file', client=client, format=format, **kwargs)
+
     def test_read(self):
         client = S3Client(AWS_ACCESS_KEY, AWS_SECRET_KEY)
         create_bucket()
@@ -81,10 +86,23 @@ def test_read(self):
         file_str = read_file.read()
         self.assertEqual(self.tempFileContents, file_str.encode('utf-8'))
 
+    def test_read_with_session(self):
+        client = S3Client(AWS_ACCESS_KEY, AWS_SECRET_KEY, AWS_SESSION_TOKEN)
+        create_bucket()
+        client.put(self.tempFilePath, 's3://mybucket/tempfile-with-session')
+        t = S3Target('s3://mybucket/tempfile-with-session', client=client)
+        read_file = t.open()
+        file_str = read_file.read()
+        self.assertEqual(self.tempFileContents, file_str.encode('utf-8'))
+
     def test_read_no_file(self):
         t = self.create_target()
         self.assertRaises(FileNotFoundException, t.open)
 
+    def test_read_no_file_with_session(self):
+        t = self.create_target_with_session()
+        self.assertRaises(FileNotFoundException, t.open)
+
     def test_read_no_file_sse(self):
         t = self.create_target(encrypt_key=True)
         self.assertRaises(FileNotFoundException, t.open)
@@ -186,6 +204,10 @@ def test_put(self):
         s3_client.put(self.tempFilePath, 's3://mybucket/putMe')
         self.assertTrue(s3_client.exists('s3://mybucket/putMe'))
 
+        s3_client = S3Client(AWS_ACCESS_KEY, AWS_SECRET_KEY, AWS_SESSION_TOKEN)
+        s3_client.put(self.tempFilePath, 's3://mybucket/putMe')
+        self.assertTrue(s3_client.exists('s3://mybucket/putMe'))
+
     def test_put_no_such_bucket(self):
         # intentionally don't create bucket
         s3_client = S3Client(AWS_ACCESS_KEY, AWS_SECRET_KEY)