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():
@@ -72,6 +73,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,14 +87,32 @@ def test_read(self):
         file_str = read_file.read()
         self.assertEqual(self.tempFileContents, file_str.encode('utf-8'))
 
+    def test_read_session(self):
+        client = S3Client(AWS_ACCESS_KEY, AWS_SECRET_KEY, AWS_SESSION_TOKEN)
+        create_bucket()
+        client.put(self.tempFilePath, 's3://mybucket/tempfile')
+        t = S3Target('s3://mybucket/tempfile', 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)
+
+        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)
 
+        t = self.create_target_with_session(encrypt_key=True)
+        self.assertRaises(FileNotFoundException, t.open)
+
     def test_read_iterator_long(self):
         # write a file that is 5X the boto buffersize
         # to test line buffering
@@ -186,6 +210,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)