Пример #1
0
class TestAWSDataSyncHookMocked(unittest.TestCase):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.source_server_hostname = 'host'
        self.source_subdirectory = 'somewhere'
        self.destination_bucket_name = 'my_bucket'
        self.destination_bucket_dir = 'dir'

        self.client = None
        self.hook = None
        self.source_location_arn = None
        self.destination_location_arn = None
        self.task_arn = None

    def setUp(self):
        self.client = boto3.client("datasync", region_name="us-east-1")
        self.hook = AWSDataSyncHook(aws_conn_id='aws_default', wait_interval_seconds=0)

        # Create default locations and tasks
        self.source_location_arn = self.client.create_location_smb(
            ServerHostname=self.source_server_hostname,
            Subdirectory=self.source_subdirectory,
            User='',
            Password='',
            AgentArns=['stuff']
        )['LocationArn']
        self.destination_location_arn = self.client.create_location_s3(
            S3BucketArn='arn:aws:s3:::{0}'.format(
                self.destination_bucket_name),
            Subdirectory=self.destination_bucket_dir,
            S3Config={'BucketAccessRoleArn': 'role'}
        )['LocationArn']
        self.task_arn = self.client.create_task(
            SourceLocationArn=self.source_location_arn,
            DestinationLocationArn=self.destination_location_arn
        )['TaskArn']

    def tearDown(self):
        # Delete all tasks:
        tasks = self.client.list_tasks()
        for task in tasks['Tasks']:
            self.client.delete_task(TaskArn=task['TaskArn'])
        # Delete all locations:
        locations = self.client.list_locations()
        for location in locations['Locations']:
            self.client.delete_location(LocationArn=location['LocationArn'])
        self.client = None

    def test_init(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        self.assertIsNone(self.hook.conn)
        self.assertFalse(self.hook.locations)
        self.assertFalse(self.hook.tasks)
        self.assertEqual(self.hook.wait_interval_seconds, 0)

    def test_create_location_smb(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        locations = self.hook.get_conn().list_locations()
        self.assertEqual(len(locations['Locations']), 2)

        server_hostname = 'my.hostname'
        subdirectory = 'my_dir'
        agent_arns = ['stuff']
        user = '******'
        domain = 'COMPANY.DOMAIN'
        mount_options = {'Version': 'SMB2'}

        location_uri = 'smb://{0}/{1}'.format(server_hostname, subdirectory)

        create_location_kwargs = {
            'ServerHostname': server_hostname,
            'Subdirectory': subdirectory,
            'User': user,
            'Password': '******',
            'Domain': domain,
            'AgentArns': agent_arns,
            'MountOptions': mount_options
        }
        location_arn = self.hook.create_location(
            location_uri, **create_location_kwargs)
        self.assertIsNotNone(location_arn)

        locations = self.client.list_locations()
        self.assertEqual(len(locations['Locations']), 3)

        location_desc = self.client.describe_location_smb(LocationArn=location_arn)
        self.assertEqual(location_desc['LocationArn'], location_arn)
        self.assertEqual(location_desc['LocationUri'], location_uri)
        self.assertEqual(location_desc['AgentArns'], agent_arns)
        self.assertEqual(location_desc['User'], user)
        self.assertEqual(location_desc['Domain'], domain)
        self.assertEqual(location_desc['MountOptions'], mount_options)

    def test_create_location_s3(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        locations = self.hook.get_conn().list_locations()
        self.assertEqual(len(locations['Locations']), 2)

        s3_bucket_arn = 'some_s3_arn'
        subdirectory = 'my_subdir'
        s3_config = {'BucketAccessRoleArn': 'myrole'}

        location_uri = 's3://{0}/{1}'.format(s3_bucket_arn, subdirectory)

        create_location_kwargs = {
            'S3BucketArn': s3_bucket_arn,
            'Subdirectory': subdirectory,
            'S3Config': s3_config
        }
        location_arn = self.hook.create_location(
            location_uri, **create_location_kwargs)
        self.assertIsNotNone(location_arn)

        locations = self.client.list_locations()
        self.assertEqual(len(locations['Locations']), 3)

        location_desc = self.client.describe_location_s3(LocationArn=location_arn)
        self.assertEqual(location_desc['LocationArn'], location_arn)
        self.assertEqual(location_desc['LocationUri'], location_uri)
        self.assertEqual(location_desc['S3Config'], s3_config)

    def test_create_task(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        log_group_arn = 'cloudwatcharn123'
        name = 'my_task'

        options = {  # Random options
            'VerifyMode': 'NONE',
            'Atime': 'NONE',
            'Mtime': 'NONE',
            'Uid': 'BOTH',
            'Gid': 'INT_VALUE',
            'PreserveDeletedFiles': 'PRESERVE',
            'PreserveDevices': 'PRESERVE',
            'PosixPermissions': 'BEST_EFFORT',
            'BytesPerSecond': 123,
        }

        create_task_kwargs = {
            'CloudWatchLogGroupArn': log_group_arn,
            'Name': name,
            'Options': options
        }

        task_arn = self.hook.create_task(
            source_location_arn=self.source_location_arn,
            destination_location_arn=self.destination_location_arn,
            **create_task_kwargs
        )

        task = self.client.describe_task(TaskArn=task_arn)
        self.assertEqual(task['TaskArn'], task_arn)
        self.assertEqual(task['Name'], name)
        self.assertEqual(task['CloudWatchLogGroupArn'], log_group_arn)
        self.assertEqual(task['Options'], options)

    def test_update_task(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_arn = self.task_arn

        task = self.client.describe_task(TaskArn=task_arn)
        self.assertNotIn('Name', task)

        update_task_kwargs = {'Name': 'xyz'}
        self.hook.update_task(task_arn, **update_task_kwargs)

        task = self.client.describe_task(TaskArn=task_arn)
        self.assertEqual(task['Name'], 'xyz')

    def test_delete_task(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_arn = self.task_arn

        tasks = self.client.list_tasks()
        self.assertEqual(len(tasks['Tasks']), 1)

        self.hook.delete_task(task_arn)

        tasks = self.client.list_tasks()
        self.assertEqual(len(tasks['Tasks']), 0)

    def test_get_location_arns(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        # Get true location_arn from boto/moto self.client
        location_uri = 'smb://{0}/{1}'.format(
            self.source_server_hostname, self.source_subdirectory)
        locations = self.client.list_locations()
        for location in locations['Locations']:
            if location['LocationUri'] == location_uri:
                location_arn = location['LocationArn']

        # Verify our self.hook gets the same
        location_arns = self.hook.get_location_arns(location_uri)

        self.assertEqual(len(location_arns), 1)
        self.assertEqual(location_arns[0], location_arn)

    def test_get_task_arns_for_location_arns(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_arns = self.hook.get_task_arns_for_location_arns(
            [self.source_location_arn],
            [self.destination_location_arn]
        )
        self.assertEqual(len(task_arns), 1)
        self.assertEqual(task_arns[0], self.task_arn)

        task_arns = self.hook.get_task_arns_for_location_arns(
            ['foo'], ['bar']
        )
        self.assertEqual(len(task_arns), 0)

    def test_start_task_execution(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task = self.client.describe_task(TaskArn=self.task_arn)
        self.assertNotIn('CurrentTaskExecutionArn', task)

        task_execution_arn = self.hook.start_task_execution(self.task_arn)
        self.assertIsNotNone(task_execution_arn)

        task = self.client.describe_task(TaskArn=self.task_arn)
        self.assertIn('CurrentTaskExecutionArn', task)
        self.assertEqual(task['CurrentTaskExecutionArn'], task_execution_arn)

        task_execution = self.client.describe_task_execution(
            TaskExecutionArn=task_execution_arn)
        self.assertIn('Status', task_execution)

    def test_cancel_task_execution(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_execution_arn = self.hook.start_task_execution(self.task_arn)
        self.assertIsNotNone(task_execution_arn)

        self.hook.cancel_task_execution(task_execution_arn=task_execution_arn)

        task = self.client.describe_task(TaskArn=self.task_arn)
        self.assertNotIn('CurrentTaskExecutionArn', task)

    def test_get_task_description(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task = self.client.describe_task(TaskArn=self.task_arn)
        self.assertIn('TaskArn', task)
        self.assertIn('Status', task)
        self.assertIn('SourceLocationArn', task)
        self.assertIn('DestinationLocationArn', task)
        self.assertNotIn('CurrentTaskExecutionArn', task)

    def test_get_current_task_execution_arn(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_execution_arn = self.hook.start_task_execution(self.task_arn)

        current_task_execution = self.hook.get_current_task_execution_arn(
            self.task_arn)
        self.assertEqual(current_task_execution, task_execution_arn)

    def test_wait_for_task_execution(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_execution_arn = self.hook.start_task_execution(self.task_arn)
        result = self.hook.wait_for_task_execution(
            task_execution_arn, max_iterations=20)

        self.assertIsNotNone(result)

    def test_wait_for_task_execution_timeout(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_execution_arn = self.hook.start_task_execution(self.task_arn)
        with self.assertRaises(AirflowTaskTimeout):
            result = self.hook.wait_for_task_execution(
                task_execution_arn, max_iterations=1)
            self.assertIsNone(result)
Пример #2
0
class TestAWSDataSyncHookMocked(unittest.TestCase):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.source_server_hostname = "host"
        self.source_subdirectory = "somewhere"
        self.destination_bucket_name = "my_bucket"
        self.destination_bucket_dir = "dir"

        self.client = None
        self.hook = None
        self.source_location_arn = None
        self.destination_location_arn = None
        self.task_arn = None

    def setUp(self):
        self.client = boto3.client("datasync", region_name="us-east-1")
        self.hook = AWSDataSyncHook(aws_conn_id="aws_default",
                                    wait_interval_seconds=0)

        # Create default locations and tasks
        self.source_location_arn = self.client.create_location_smb(
            ServerHostname=self.source_server_hostname,
            Subdirectory=self.source_subdirectory,
            User="",
            Password="",
            AgentArns=["stuff"],
        )["LocationArn"]
        self.destination_location_arn = self.client.create_location_s3(
            S3BucketArn="arn:aws:s3:::{0}".format(
                self.destination_bucket_name),
            Subdirectory=self.destination_bucket_dir,
            S3Config={"BucketAccessRoleArn": "role"},
        )["LocationArn"]
        self.task_arn = self.client.create_task(
            SourceLocationArn=self.source_location_arn,
            DestinationLocationArn=self.destination_location_arn,
        )["TaskArn"]

    def tearDown(self):
        # Delete all tasks:
        tasks = self.client.list_tasks()
        for task in tasks["Tasks"]:
            self.client.delete_task(TaskArn=task["TaskArn"])
        # Delete all locations:
        locations = self.client.list_locations()
        for location in locations["Locations"]:
            self.client.delete_location(LocationArn=location["LocationArn"])
        self.client = None

    def test_init(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        self.assertIsNone(self.hook.conn)
        self.assertFalse(self.hook.locations)
        self.assertFalse(self.hook.tasks)
        self.assertEqual(self.hook.wait_interval_seconds, 0)

    def test_create_location_smb(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        locations = self.hook.get_conn().list_locations()
        self.assertEqual(len(locations["Locations"]), 2)

        server_hostname = "my.hostname"
        subdirectory = "my_dir"
        agent_arns = ["stuff"]
        user = "******"
        domain = "COMPANY.DOMAIN"
        mount_options = {"Version": "SMB2"}

        location_uri = "smb://{0}/{1}".format(server_hostname, subdirectory)

        create_location_kwargs = {
            "ServerHostname": server_hostname,
            "Subdirectory": subdirectory,
            "User": user,
            "Password": "******",
            "Domain": domain,
            "AgentArns": agent_arns,
            "MountOptions": mount_options,
        }
        location_arn = self.hook.create_location(location_uri,
                                                 **create_location_kwargs)
        self.assertIsNotNone(location_arn)

        locations = self.client.list_locations()
        self.assertEqual(len(locations["Locations"]), 3)

        location_desc = self.client.describe_location_smb(
            LocationArn=location_arn)
        self.assertEqual(location_desc["LocationArn"], location_arn)
        self.assertEqual(location_desc["LocationUri"], location_uri)
        self.assertEqual(location_desc["AgentArns"], agent_arns)
        self.assertEqual(location_desc["User"], user)
        self.assertEqual(location_desc["Domain"], domain)
        self.assertEqual(location_desc["MountOptions"], mount_options)

    def test_create_location_s3(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        locations = self.hook.get_conn().list_locations()
        self.assertEqual(len(locations["Locations"]), 2)

        s3_bucket_arn = "some_s3_arn"
        subdirectory = "my_subdir"
        s3_config = {"BucketAccessRoleArn": "myrole"}

        location_uri = "s3://{0}/{1}".format(s3_bucket_arn, subdirectory)

        create_location_kwargs = {
            "S3BucketArn": s3_bucket_arn,
            "Subdirectory": subdirectory,
            "S3Config": s3_config,
        }
        location_arn = self.hook.create_location(location_uri,
                                                 **create_location_kwargs)
        self.assertIsNotNone(location_arn)

        locations = self.client.list_locations()
        self.assertEqual(len(locations["Locations"]), 3)

        location_desc = self.client.describe_location_s3(
            LocationArn=location_arn)
        self.assertEqual(location_desc["LocationArn"], location_arn)
        self.assertEqual(location_desc["LocationUri"], location_uri)
        self.assertEqual(location_desc["S3Config"], s3_config)

    def test_create_task(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        log_group_arn = "cloudwatcharn123"
        name = "my_task"

        options = {  # Random options
            "VerifyMode": "NONE",
            "Atime": "NONE",
            "Mtime": "NONE",
            "Uid": "BOTH",
            "Gid": "INT_VALUE",
            "PreserveDeletedFiles": "PRESERVE",
            "PreserveDevices": "PRESERVE",
            "PosixPermissions": "BEST_EFFORT",
            "BytesPerSecond": 123,
        }

        create_task_kwargs = {
            "CloudWatchLogGroupArn": log_group_arn,
            "Name": name,
            "Options": options,
        }

        task_arn = self.hook.create_task(
            source_location_arn=self.source_location_arn,
            destination_location_arn=self.destination_location_arn,
            **create_task_kwargs)

        task = self.client.describe_task(TaskArn=task_arn)
        self.assertEqual(task["TaskArn"], task_arn)
        self.assertEqual(task["Name"], name)
        self.assertEqual(task["CloudWatchLogGroupArn"], log_group_arn)
        self.assertEqual(task["Options"], options)

    def test_update_task(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_arn = self.task_arn

        task = self.client.describe_task(TaskArn=task_arn)
        self.assertNotIn("Name", task)

        update_task_kwargs = {"Name": "xyz"}
        self.hook.update_task(task_arn, **update_task_kwargs)

        task = self.client.describe_task(TaskArn=task_arn)
        self.assertEqual(task["Name"], "xyz")

    def test_delete_task(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_arn = self.task_arn

        tasks = self.client.list_tasks()
        self.assertEqual(len(tasks["Tasks"]), 1)

        self.hook.delete_task(task_arn)

        tasks = self.client.list_tasks()
        self.assertEqual(len(tasks["Tasks"]), 0)

    def test_get_location_arns(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        # Get true location_arn from boto/moto self.client
        location_uri = "smb://{0}/{1}".format(self.source_server_hostname,
                                              self.source_subdirectory)
        locations = self.client.list_locations()
        for location in locations["Locations"]:
            if location["LocationUri"] == location_uri:
                location_arn = location["LocationArn"]

        # Verify our self.hook gets the same
        location_arns = self.hook.get_location_arns(location_uri)

        self.assertEqual(len(location_arns), 1)
        self.assertEqual(location_arns[0], location_arn)

    def test_get_location_arns_case_sensitive(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        # Get true location_arn from boto/moto self.client
        location_uri = "smb://{0}/{1}".format(
            self.source_server_hostname.upper(), self.source_subdirectory)
        locations = self.client.list_locations()
        for location in locations["Locations"]:
            if location["LocationUri"] == location_uri.lower():
                location_arn = location["LocationArn"]

        # Verify our self.hook can do case sensitive searches
        location_arns = self.hook.get_location_arns(location_uri,
                                                    case_sensitive=True)
        self.assertEqual(len(location_arns), 0)
        location_arns = self.hook.get_location_arns(location_uri,
                                                    case_sensitive=False)
        self.assertEqual(len(location_arns), 1)
        self.assertEqual(location_arns[0], location_arn)

    def test_get_location_arns_trailing_slash(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        # Get true location_arn from boto/moto self.client
        location_uri = "smb://{0}/{1}/".format(self.source_server_hostname,
                                               self.source_subdirectory)
        locations = self.client.list_locations()
        for location in locations["Locations"]:
            if location["LocationUri"] == location_uri[:-1]:
                location_arn = location["LocationArn"]

        # Verify our self.hook manages trailing / correctly
        location_arns = self.hook.get_location_arns(
            location_uri, ignore_trailing_slash=False)
        self.assertEqual(len(location_arns), 0)
        location_arns = self.hook.get_location_arns(location_uri,
                                                    ignore_trailing_slash=True)
        self.assertEqual(len(location_arns), 1)
        self.assertEqual(location_arns[0], location_arn)

    def test_get_task_arns_for_location_arns(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_arns = self.hook.get_task_arns_for_location_arns(
            [self.source_location_arn], [self.destination_location_arn])
        self.assertEqual(len(task_arns), 1)
        self.assertEqual(task_arns[0], self.task_arn)

        task_arns = self.hook.get_task_arns_for_location_arns(["foo"], ["bar"])
        self.assertEqual(len(task_arns), 0)

    def test_start_task_execution(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task = self.client.describe_task(TaskArn=self.task_arn)
        self.assertNotIn("CurrentTaskExecutionArn", task)

        task_execution_arn = self.hook.start_task_execution(self.task_arn)
        self.assertIsNotNone(task_execution_arn)

        task = self.client.describe_task(TaskArn=self.task_arn)
        self.assertIn("CurrentTaskExecutionArn", task)
        self.assertEqual(task["CurrentTaskExecutionArn"], task_execution_arn)

        task_execution = self.client.describe_task_execution(
            TaskExecutionArn=task_execution_arn)
        self.assertIn("Status", task_execution)

    def test_cancel_task_execution(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_execution_arn = self.hook.start_task_execution(self.task_arn)
        self.assertIsNotNone(task_execution_arn)

        self.hook.cancel_task_execution(task_execution_arn=task_execution_arn)

        task = self.client.describe_task(TaskArn=self.task_arn)
        self.assertNotIn("CurrentTaskExecutionArn", task)

    def test_get_task_description(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task = self.client.describe_task(TaskArn=self.task_arn)
        self.assertIn("TaskArn", task)
        self.assertIn("Status", task)
        self.assertIn("SourceLocationArn", task)
        self.assertIn("DestinationLocationArn", task)
        self.assertNotIn("CurrentTaskExecutionArn", task)

    def test_get_current_task_execution_arn(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_execution_arn = self.hook.start_task_execution(self.task_arn)

        current_task_execution = self.hook.get_current_task_execution_arn(
            self.task_arn)
        self.assertEqual(current_task_execution, task_execution_arn)

    def test_wait_for_task_execution(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_execution_arn = self.hook.start_task_execution(self.task_arn)
        result = self.hook.wait_for_task_execution(task_execution_arn,
                                                   max_iterations=20)

        self.assertIsNotNone(result)

    def test_wait_for_task_execution_timeout(self, mock_get_conn):
        # ### Configure mock:
        mock_get_conn.return_value = self.client
        # ### Begin tests:

        task_execution_arn = self.hook.start_task_execution(self.task_arn)
        with self.assertRaises(AirflowTaskTimeout):
            result = self.hook.wait_for_task_execution(task_execution_arn,
                                                       max_iterations=1)
            self.assertIsNone(result)