updating functions, adding unit tests

core/eolearn/core/eodata.py

@@ -559,13 +559,15 @@ def save(self, path, features=..., file_format=FileFormat.NPY,
                 shutil.rmtree(tmp_path)
             raise ex
 
-    def save_aws(self, bucket_name, patch_location):
+    def save_aws(self, bucket_name, patch_location, s3client=None):
         """Saves EOPatch to the AWS S3 bucket. AWS credentials should be properly configured.
 
         :param bucket_name: Name of the AWS S3 bucket
         :type bucket_name: str
         :param patch_location: Location of the EOPatch on the AWS S3 bucket
         :type patch_location: str
+        :param s3client: Override the automatic s3 client
+        :type s3client: botocore.client.S3
         """
         features = [feature for feature in self.get_features() if not feature.is_meta()]
         features = [(ftype, fname) for ftype in features for fname in self[ftype].keys()]
@@ -580,7 +582,7 @@ def save_aws(self, bucket_name, patch_location):
 
         streams = (pickle.dumps(self[feat], protocol=pickle.HIGHEST_PROTOCOL) for feat in features)
 
-        s3client = boto3.client('s3')
+        s3client = boto3.client('s3') if s3client is None else s3client
         for stream, path in zip(streams, paths):
             s3client.put_object(Bucket=bucket_name, Key=path, Body=stream)
 
@@ -716,15 +718,19 @@ def load(path, features=..., lazy_loading=False, mmap=False):
         return EOPatch(**requested_content)
 
     @staticmethod
-    def load_aws(bucket_name, patch_location):
+    def load_aws(bucket_name, patch_location, s3client=None):
         """Loads EOPatch from the AWS S3 bucket. AWS credentials should be properly configured.
 
         :param bucket_name: Name of the AWS S3 bucket
         :type bucket_name: str
         :param patch_location: Location of the EOPatch on the AWS S3 bucket
         :type patch_location: str
+        :param s3client: Override the automatic s3 client
+        :type s3client: botocore.client.S3
         """
-        s3client = boto3.client('s3')
+
+        s3client = boto3.client('s3') if s3client is None else s3client
+
         patch_location += '/' if not patch_location.endswith('/') else ''
         list_request = s3client.list_objects(Bucket=bucket_name, Prefix=patch_location)
 

core/eolearn/tests/test_aws.py

@@ -0,0 +1,71 @@
+"""
+Credits:
+Copyright (c) 2017-2019 Matej Aleksandrov, Matej Batič, Andrej Burja, Eva Erzin (Sinergise)
+Copyright (c) 2017-2019 Grega Milčinski, Matic Lubej, Devis Peresutti, Jernej Puc, Tomislav Slijepčević (Sinergise)
+Copyright (c) 2017-2019 Blaž Sovdat, Jovan Višnjić, Anže Zupanc, Lojze Žust (Sinergise)
+
+This source code is licensed under the MIT license found in the LICENSE
+file in the root directory of this source tree.
+"""
+
+import unittest
+import logging
+import os
+import pickle
+import boto3
+from moto import mock_s3
+
+from eolearn.core import EOPatch
+
+logging.basicConfig(level=logging.DEBUG)
+
+
+class TestEOPatchAWS(unittest.TestCase):
+    PATCH_FILENAME = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../../../example_data/TestEOPatchAWS')
+    eop = EOPatch.load(PATCH_FILENAME)
+
+    MY_BUCKET = "test-bucket"
+    MY_PREFIX = "TestEOPatchAWS"
+
+    @mock_s3
+    def create_and_fill_s3_bucket(self):
+        s3client = boto3.client('s3', region_name='eu-central-1')
+        s3resource = boto3.resource('s3', region_name='eu-central-1')
+        s3resource.create_bucket(Bucket=self.MY_BUCKET)
+        self.eop.save_aws(bucket_name=self.MY_BUCKET, patch_location=self.MY_PREFIX, s3client=s3client)
+        return s3client, s3resource
+
+    def empty_and_delete_s3_bucket(self, s3resource):
+        bucket = s3resource.Bucket(self.MY_BUCKET)
+        for key in bucket.objects.all():
+            key.delete()
+        bucket.delete()
+
+    @mock_s3
+    def test_save_eopatch(self):
+        stats_filename = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'stats/test_save_stats.pkl')
+        saved_content = pickle.load(open(stats_filename, 'rb'))
+
+        s3client, s3resource = self.create_and_fill_s3_bucket()
+
+        response = s3client.list_objects(Bucket=self.MY_BUCKET, Prefix=self.MY_PREFIX)
+        content = [x['Key'] for x in response['Contents']]
+
+        update_stats = True
+        if update_stats:
+            pickle.dump(content, open(stats_filename, 'wb'))
+
+        self.empty_and_delete_s3_bucket(s3resource)
+        self.assertEqual(saved_content, content)
+
+    @mock_s3
+    def test_load_eopatch(self):
+        s3client, s3resource = self.create_and_fill_s3_bucket()
+        eop = EOPatch.load_aws(bucket_name=self.MY_BUCKET, patch_location=self.MY_PREFIX, s3client=s3client)
+
+        self.empty_and_delete_s3_bucket(s3resource)
+        self.assertEqual(self.eop, eop)
+
+
+if __name__ == '__main__':
+    unittest.main()

core/requirements.txt

@@ -4,3 +4,5 @@ numpy>=1.16
 geopandas
 sentinelhub>=2.5.0
 tqdm>=4.27
+boto3
+moto