s3_file_uploader_test_base.py
Go to the documentation of this file.
1 #!/usr/bin/env python
2 # Copyright (c) 2020, Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License").
5 # You may not use this file except in compliance with the License.
6 # A copy of the License is located at
7 #
8 # http://aws.amazon.com/apache2.0
9 #
10 # or in the "license" file accompanying this file. This file is distributed
11 # on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12 # express or implied. See the License for the specific language governing
13 # permissions and limitations under the License.
14 
15 import filecmp
16 import os
17 import tempfile
18 import unittest
19 
20 import actionlib
21 
22 from file_uploader_msgs.msg import UploadFilesAction, UploadFilesGoal
23 
24 import rospy
25 
26 import rostest
27 
28 from s3_client import S3Client
29 
30 ACTION = '/s3_file_uploader/UploadFiles'
31 TEST_NODE_NAME = 'upload_files_action_client'
32 AWS_DEFAULT_REGION = 'us-west-2'
33 
34 
35 class S3FileUploaderTestBase(unittest.TestCase):
36  @classmethod
37  def setUpClass(cls):
38  rospy.init_node(TEST_NODE_NAME, log_level=rospy.DEBUG)
39  s3 = S3Client(cls.extract_s3_region())
40  s3_bucket_name = rospy.get_param('/s3_file_uploader/s3_bucket')
41  s3.create_bucket(s3_bucket_name)
42  s3.wait_for_bucket_create(s3_bucket_name)
43 
44  @classmethod
45  def tearDownClass(cls):
46  s3 = S3Client(cls.extract_s3_region())
47  s3_bucket_name = rospy.get_param('/s3_file_uploader/s3_bucket')
48  s3.delete_bucket(s3_bucket_name)
49 
50  @classmethod
52  s3_region = rospy.get_param(
53  '/s3_file_uploader/aws_client_configuration/region')
54  if not s3_region:
55  return AWS_DEFAULT_REGION
56  return s3_region
57 
58  def setUp(self):
59  self.action_client = None
60  self.s3_bucket = rospy.get_param('/s3_file_uploader/s3_bucket')
61  self.s3_region = S3FileUploaderTestBase.extract_s3_region()
63  self.s3_key_prefix = 'foo/bar'
65  self.files_to_delete = []
66 
67  def tearDown(self):
68  # delete any temp files created on disk
69  for file_name in self.files_to_delete:
70  os.remove(file_name)
71  # delete any temp files created in the bucket as a part of tests
72  contents = self.s3_client.list_objects(self.s3_bucket)
73  objects_to_delete = map(
74  lambda content: content['Key'],
75  filter(
76  lambda obj: obj['Key'] in self.objects_to_delete,
77  contents
78  )
79  )
80  if len(objects_to_delete) > 0:
81  self.s3_client.delete_objects(self.s3_bucket, objects_to_delete)
82 
84  self.action_client = actionlib.SimpleActionClient(ACTION, UploadFilesAction)
85  res = self.action_client.wait_for_server()
86  self.assertTrue(res, 'Failed to connect to action server')
87 
88  def _upload_temp_files(self, temp_file_names):
89  goal = UploadFilesGoal(
90  upload_location=self.s3_key_prefix,
91  files=temp_file_names
92  )
93  self.action_client.send_goal(goal)
94  self.action_client.wait_for_result(rospy.Duration.from_sec(15.0))
95  return self.action_client.get_result()
96 
97  def _assert_successful_upload(self, result, temp_file_names):
98  self.assertEqual(
99  len(result.files_uploaded),
100  len(temp_file_names),
101  'Found %d files' % len(result.files_uploaded)
102  )
103  for uploaded_file_name, temp_file_name in zip(result.files_uploaded, temp_file_names):
105  uploaded_file_name,
106  temp_file_name
107  )
108 
109  def _assert_successful_file_upload(self, uploaded_s3_file_path, temp_file_name):
110  # mark the uploaded test file for delete upon tearDown
111  self.objects_to_delete.append(uploaded_s3_file_path)
112  uploaded_s3_key_prefix, uploaded_s3_key = os.path.split(
113  uploaded_s3_file_path)
114  self.assertEqual(
115  uploaded_s3_key_prefix,
116  self.s3_key_prefix,
117  'File uploaded to unexpected location in s3')
118  self.assertEqual(
119  uploaded_s3_key,
120  os.path.basename(temp_file_name),
121  'Uploaded file name mismatch')
122  self._assert_uploaded_content(uploaded_s3_file_path, temp_file_name)
123 
124  def _assert_uploaded_content(self, s3_key, temp_file_name):
125  with tempfile.NamedTemporaryFile() as f:
126  self.s3_client.download_file(self.s3_bucket, s3_key, f.name)
127  self.assertTrue(
128  filecmp.cmp(f.name, temp_file_name),
129  'Local file content did not match uploaded content')
def _assert_successful_file_upload(self, uploaded_s3_file_path, temp_file_name)
def _assert_successful_upload(self, result, temp_file_names)


integ_tests
Author(s): AWS RoboMaker
autogenerated on Thu Mar 4 2021 03:52:30