How to use mock_session method in tavern

Best Python code snippet using tavern

test_sagemaker_hook.py

Source:test_sagemaker_hook.py

1# -*- coding: utf-8 -*-2#3# Licensed to the Apache Software Foundation (ASF) under one4# or more contributor license agreements.  See the NOTICE file5# distributed with this work for additional information6# regarding copyright ownership.  The ASF licenses this file7# to you under the Apache License, Version 2.0 (the8# "License"); you may not use this file except in compliance9# with the License.  You may obtain a copy of the License at10#11#   http://www.apache.org/licenses/LICENSE-2.012#13# Unless required by applicable law or agreed to in writing,14# software distributed under the License is distributed on an15# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY16# KIND, either express or implied.  See the License for the17# specific language governing permissions and limitations18# under the License.19#20import unittest21import time22from datetime import datetime23from tzlocal import get_localzone24try:25    from unittest import mock26except ImportError:27    try:28        import mock29    except ImportError:30        mock = None31from airflow import configuration32from airflow.contrib.hooks.sagemaker_hook import (SageMakerHook, secondary_training_status_changed,33                                                  secondary_training_status_message, LogState)34from airflow.hooks.S3_hook import S3Hook35from airflow.exceptions import AirflowException36role = 'arn:aws:iam:role/test-role'37path = 'local/data'38bucket = 'test-bucket'39key = 'test/data'40data_url = 's3://{}/{}'.format(bucket, key)41job_name = 'test-job'42model_name = 'test-model'43config_name = 'test-endpoint-config'44endpoint_name = 'test-endpoint'45image = 'test-image'46test_arn_return = {'Arn': 'testarn'}47output_url = 's3://{}/test/output'.format(bucket)48create_training_params = {49    'AlgorithmSpecification': {50        'TrainingImage': image,51        'TrainingInputMode': 'File'52    },53    'RoleArn': role,54    'OutputDataConfig': {55        'S3OutputPath': output_url56    },57    'ResourceConfig': {58        'InstanceCount': 2,59        'InstanceType': 'ml.c4.8xlarge',60        'VolumeSizeInGB': 5061    },62    'TrainingJobName': job_name,63    'HyperParameters': {64        'k': '10',65        'feature_dim': '784',66        'mini_batch_size': '500',67        'force_dense': 'True'68    },69    'StoppingCondition': {70        'MaxRuntimeInSeconds': 60 * 6071    },72    'InputDataConfig': [73        {74            'ChannelName': 'train',75            'DataSource': {76                'S3DataSource': {77                    'S3DataType': 'S3Prefix',78                    'S3Uri': data_url,79                    'S3DataDistributionType': 'FullyReplicated'80                }81            },82            'CompressionType': 'None',83            'RecordWrapperType': 'None'84        }85    ]86}87create_tuning_params = {88    'HyperParameterTuningJobName': job_name,89    'HyperParameterTuningJobConfig': {90        'Strategy': 'Bayesian',91        'HyperParameterTuningJobObjective': {92            'Type': 'Maximize',93            'MetricName': 'test_metric'94        },95        'ResourceLimits': {96            'MaxNumberOfTrainingJobs': 123,97            'MaxParallelTrainingJobs': 12398        },99        'ParameterRanges': {100            'IntegerParameterRanges': [101                {102                    'Name': 'k',103                    'MinValue': '2',104                    'MaxValue': '10'105                },106            ]107        }108    },109    'TrainingJobDefinition': {110        'StaticHyperParameters': create_training_params['HyperParameters'],111        'AlgorithmSpecification': create_training_params['AlgorithmSpecification'],112        'RoleArn': 'string',113        'InputDataConfig': create_training_params['InputDataConfig'],114        'OutputDataConfig': create_training_params['OutputDataConfig'],115        'ResourceConfig': create_training_params['ResourceConfig'],116        'StoppingCondition': dict(MaxRuntimeInSeconds=60 * 60)117    }118}119create_transform_params = {120    'TransformJobName': job_name,121    'ModelName': model_name,122    'BatchStrategy': 'MultiRecord',123    'TransformInput': {124        'DataSource': {125            'S3DataSource': {126                'S3DataType': 'S3Prefix',127                'S3Uri': data_url128            }129        }130    },131    'TransformOutput': {132        'S3OutputPath': output_url,133    },134    'TransformResources': {135        'InstanceType': 'ml.m4.xlarge',136        'InstanceCount': 123137    }138}139create_model_params = {140    'ModelName': model_name,141    'PrimaryContainer': {142        'Image': image,143        'ModelDataUrl': output_url,144    },145    'ExecutionRoleArn': role146}147create_endpoint_config_params = {148    'EndpointConfigName': config_name,149    'ProductionVariants': [150        {151            'VariantName': 'AllTraffic',152            'ModelName': model_name,153            'InitialInstanceCount': 1,154            'InstanceType': 'ml.c4.xlarge'155        }156    ]157}158create_endpoint_params = {159    'EndpointName': endpoint_name,160    'EndpointConfigName': config_name161}162update_endpoint_params = create_endpoint_params163DESCRIBE_TRAINING_COMPELETED_RETURN = {164    'TrainingJobStatus': 'Completed',165    'ResourceConfig': {166        'InstanceCount': 1,167        'InstanceType': 'ml.c4.xlarge',168        'VolumeSizeInGB': 10169    },170    'TrainingStartTime': datetime(2018, 2, 17, 7, 15, 0, 103000),171    'TrainingEndTime': datetime(2018, 2, 17, 7, 19, 34, 953000),172    'ResponseMetadata': {173        'HTTPStatusCode': 200,174    }175}176DESCRIBE_TRAINING_INPROGRESS_RETURN = dict(DESCRIBE_TRAINING_COMPELETED_RETURN)177DESCRIBE_TRAINING_INPROGRESS_RETURN.update({'TrainingJobStatus': 'InProgress'})178DESCRIBE_TRAINING_FAILED_RETURN = dict(DESCRIBE_TRAINING_COMPELETED_RETURN)179DESCRIBE_TRAINING_FAILED_RETURN.update({'TrainingJobStatus': 'Failed',180                                        'FailureReason': 'Unknown'})181DESCRIBE_TRAINING_STOPPING_RETURN = dict(DESCRIBE_TRAINING_COMPELETED_RETURN)182DESCRIBE_TRAINING_STOPPING_RETURN.update({'TrainingJobStatus': 'Stopping'})183message = 'message'184status = 'status'185SECONDARY_STATUS_DESCRIPTION_1 = {186    'SecondaryStatusTransitions': [{'StatusMessage': message, 'Status': status}]187}188SECONDARY_STATUS_DESCRIPTION_2 = {189    'SecondaryStatusTransitions': [{'StatusMessage': 'different message', 'Status': status}]190}191DEFAULT_LOG_STREAMS = {'logStreams': [{'logStreamName': job_name + '/xxxxxxxxx'}]}192LIFECYCLE_LOG_STREAMS = [DEFAULT_LOG_STREAMS,193                         DEFAULT_LOG_STREAMS,194                         DEFAULT_LOG_STREAMS,195                         DEFAULT_LOG_STREAMS,196                         DEFAULT_LOG_STREAMS,197                         DEFAULT_LOG_STREAMS]198DEFAULT_LOG_EVENTS = [{'nextForwardToken': None, 'events': [{'timestamp': 1, 'message': 'hi there #1'}]},199                      {'nextForwardToken': None, 'events': []}]200STREAM_LOG_EVENTS = [{'nextForwardToken': None, 'events': [{'timestamp': 1, 'message': 'hi there #1'}]},201                     {'nextForwardToken': None, 'events': []},202                     {'nextForwardToken': None, 'events': [{'timestamp': 1, 'message': 'hi there #1'},203                                                           {'timestamp': 2, 'message': 'hi there #2'}]},204                     {'nextForwardToken': None, 'events': []},205                     {'nextForwardToken': None, 'events': [{'timestamp': 2, 'message': 'hi there #2'},206                                                           {'timestamp': 2, 'message': 'hi there #2a'},207                                                           {'timestamp': 3, 'message': 'hi there #3'}]},208                     {'nextForwardToken': None, 'events': []}]209test_evaluation_config = {210    'Image': image,211    'Role': role,212    'S3Operations': {213        'S3CreateBucket': [214            {215                'Bucket': bucket216            }217        ],218        'S3Upload': [219            {220                'Path': path,221                'Bucket': bucket,222                'Key': key,223                'Tar': False224            }225        ]226    }227}228class TestSageMakerHook(unittest.TestCase):229    def setUp(self):230        configuration.load_test_config()231    @mock.patch.object(S3Hook, 'create_bucket')232    @mock.patch.object(S3Hook, 'load_file')233    def test_configure_s3_resources(self, mock_load_file, mock_create_bucket):234        hook = SageMakerHook()235        evaluation_result = {236            'Image': image,237            'Role': role238        }239        hook.configure_s3_resources(test_evaluation_config)240        self.assertEqual(test_evaluation_config, evaluation_result)241        mock_create_bucket.assert_called_once_with(bucket_name=bucket)242        mock_load_file.assert_called_once_with(path, key, bucket)243    @mock.patch.object(SageMakerHook, 'get_conn')244    @mock.patch.object(S3Hook, 'check_for_key')245    @mock.patch.object(S3Hook, 'check_for_bucket')246    @mock.patch.object(S3Hook, 'check_for_prefix')247    def test_check_s3_url(self,248                          mock_check_prefix,249                          mock_check_bucket,250                          mock_check_key,251                          mock_client):252        mock_client.return_value = None253        hook = SageMakerHook()254        mock_check_bucket.side_effect = [False, True, True, True]255        mock_check_key.side_effect = [False, True, False]256        mock_check_prefix.side_effect = [False, True, True]257        self.assertRaises(AirflowException,258                          hook.check_s3_url, data_url)259        self.assertRaises(AirflowException,260                          hook.check_s3_url, data_url)261        self.assertEqual(hook.check_s3_url(data_url), True)262        self.assertEqual(hook.check_s3_url(data_url), True)263    @mock.patch.object(SageMakerHook, 'get_conn')264    @mock.patch.object(SageMakerHook, 'check_s3_url')265    def test_check_valid_training(self, mock_check_url, mock_client):266        mock_client.return_value = None267        hook = SageMakerHook()268        hook.check_training_config(create_training_params)269        mock_check_url.assert_called_once_with(data_url)270    @mock.patch.object(SageMakerHook, 'get_conn')271    @mock.patch.object(SageMakerHook, 'check_s3_url')272    def test_check_valid_tuning(self, mock_check_url, mock_client):273        mock_client.return_value = None274        hook = SageMakerHook()275        hook.check_tuning_config(create_tuning_params)276        mock_check_url.assert_called_once_with(data_url)277    @mock.patch.object(SageMakerHook, 'get_client_type')278    def test_conn(self, mock_get_client):279        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')280        self.assertEqual(hook.aws_conn_id, 'sagemaker_test_conn_id')281    @mock.patch.object(SageMakerHook, 'check_training_config')282    @mock.patch.object(SageMakerHook, 'get_conn')283    def test_create_training_job(self, mock_client, mock_check_training):284        mock_check_training.return_value = True285        mock_session = mock.Mock()286        attrs = {'create_training_job.return_value':287                 test_arn_return}288        mock_session.configure_mock(**attrs)289        mock_client.return_value = mock_session290        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')291        response = hook.create_training_job(create_training_params,292                                            wait_for_completion=False,293                                            print_log=False)294        mock_session.create_training_job.assert_called_once_with(**create_training_params)295        self.assertEqual(response, test_arn_return)296    @mock.patch.object(SageMakerHook, 'check_training_config')297    @mock.patch.object(SageMakerHook, 'get_conn')298    def test_training_ends_with_wait(self, mock_client, mock_check_training):299        mock_check_training.return_value = True300        mock_session = mock.Mock()301        attrs = {'create_training_job.return_value':302                 test_arn_return,303                 'describe_training_job.side_effect':304                 [DESCRIBE_TRAINING_INPROGRESS_RETURN,305                  DESCRIBE_TRAINING_STOPPING_RETURN,306                  DESCRIBE_TRAINING_COMPELETED_RETURN,307                  DESCRIBE_TRAINING_COMPELETED_RETURN]308                 }309        mock_session.configure_mock(**attrs)310        mock_client.return_value = mock_session311        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id_1')312        hook.create_training_job(create_training_params, wait_for_completion=True,313                                 print_log=False, check_interval=1)314        self.assertEqual(mock_session.describe_training_job.call_count, 4)315    @mock.patch.object(SageMakerHook, 'check_training_config')316    @mock.patch.object(SageMakerHook, 'get_conn')317    def test_training_throws_error_when_failed_with_wait(318            self, mock_client, mock_check_training):319        mock_check_training.return_value = True320        mock_session = mock.Mock()321        attrs = {'create_training_job.return_value':322                 test_arn_return,323                 'describe_training_job.side_effect':324                 [DESCRIBE_TRAINING_INPROGRESS_RETURN,325                  DESCRIBE_TRAINING_STOPPING_RETURN,326                  DESCRIBE_TRAINING_FAILED_RETURN,327                  DESCRIBE_TRAINING_COMPELETED_RETURN]328                 }329        mock_session.configure_mock(**attrs)330        mock_client.return_value = mock_session331        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id_1')332        self.assertRaises(AirflowException, hook.create_training_job,333                          create_training_params, wait_for_completion=True,334                          print_log=False, check_interval=1)335        self.assertEqual(mock_session.describe_training_job.call_count, 3)336    @mock.patch.object(SageMakerHook, 'check_tuning_config')337    @mock.patch.object(SageMakerHook, 'get_conn')338    def test_create_tuning_job(self, mock_client, mock_check_tuning):339        mock_session = mock.Mock()340        attrs = {'create_hyper_parameter_tuning_job.return_value':341                 test_arn_return}342        mock_session.configure_mock(**attrs)343        mock_client.return_value = mock_session344        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')345        response = hook.create_tuning_job(create_tuning_params,346                                          wait_for_completion=False)347        mock_session.create_hyper_parameter_tuning_job.\348            assert_called_once_with(**create_tuning_params)349        self.assertEqual(response, test_arn_return)350    @mock.patch.object(SageMakerHook, 'check_s3_url')351    @mock.patch.object(SageMakerHook, 'get_conn')352    def test_create_transform_job(self, mock_client, mock_check_url):353        mock_check_url.return_value = True354        mock_session = mock.Mock()355        attrs = {'create_transform_job.return_value':356                 test_arn_return}357        mock_session.configure_mock(**attrs)358        mock_client.return_value = mock_session359        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')360        response = hook.create_transform_job(create_transform_params,361                                             wait_for_completion=False)362        mock_session.create_transform_job.assert_called_once_with(363            **create_transform_params)364        self.assertEqual(response, test_arn_return)365    @mock.patch.object(SageMakerHook, 'get_conn')366    def test_create_model(self, mock_client):367        mock_session = mock.Mock()368        attrs = {'create_model.return_value':369                 test_arn_return}370        mock_session.configure_mock(**attrs)371        mock_client.return_value = mock_session372        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')373        response = hook.create_model(create_model_params)374        mock_session.create_model.assert_called_once_with(**create_model_params)375        self.assertEqual(response, test_arn_return)376    @mock.patch.object(SageMakerHook, 'get_conn')377    def test_create_endpoint_config(self, mock_client):378        mock_session = mock.Mock()379        attrs = {'create_endpoint_config.return_value':380                 test_arn_return}381        mock_session.configure_mock(**attrs)382        mock_client.return_value = mock_session383        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')384        response = hook.create_endpoint_config(create_endpoint_config_params)385        mock_session.create_endpoint_config\386            .assert_called_once_with(**create_endpoint_config_params)387        self.assertEqual(response, test_arn_return)388    @mock.patch.object(SageMakerHook, 'get_conn')389    def test_create_endpoint(self, mock_client):390        mock_session = mock.Mock()391        attrs = {'create_endpoint.return_value':392                 test_arn_return}393        mock_session.configure_mock(**attrs)394        mock_client.return_value = mock_session395        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')396        response = hook.create_endpoint(create_endpoint_params,397                                        wait_for_completion=False)398        mock_session.create_endpoint\399            .assert_called_once_with(**create_endpoint_params)400        self.assertEqual(response, test_arn_return)401    @mock.patch.object(SageMakerHook, 'get_conn')402    def test_update_endpoint(self, mock_client):403        mock_session = mock.Mock()404        attrs = {'update_endpoint.return_value':405                 test_arn_return}406        mock_session.configure_mock(**attrs)407        mock_client.return_value = mock_session408        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')409        response = hook.update_endpoint(update_endpoint_params,410                                        wait_for_completion=False)411        mock_session.update_endpoint\412            .assert_called_once_with(**update_endpoint_params)413        self.assertEqual(response, test_arn_return)414    @mock.patch.object(SageMakerHook, 'get_conn')415    def test_describe_training_job(self, mock_client):416        mock_session = mock.Mock()417        attrs = {'describe_training_job.return_value': 'InProgress'}418        mock_session.configure_mock(**attrs)419        mock_client.return_value = mock_session420        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')421        response = hook.describe_training_job(job_name)422        mock_session.describe_training_job.\423            assert_called_once_with(TrainingJobName=job_name)424        self.assertEqual(response, 'InProgress')425    @mock.patch.object(SageMakerHook, 'get_conn')426    def test_describe_tuning_job(self, mock_client):427        mock_session = mock.Mock()428        attrs = {'describe_hyper_parameter_tuning_job.return_value':429                 'InProgress'}430        mock_session.configure_mock(**attrs)431        mock_client.return_value = mock_session432        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')433        response = hook.describe_tuning_job(job_name)434        mock_session.describe_hyper_parameter_tuning_job.\435            assert_called_once_with(HyperParameterTuningJobName=job_name)436        self.assertEqual(response, 'InProgress')437    @mock.patch.object(SageMakerHook, 'get_conn')438    def test_describe_transform_job(self, mock_client):439        mock_session = mock.Mock()440        attrs = {'describe_transform_job.return_value':441                 'InProgress'}442        mock_session.configure_mock(**attrs)443        mock_client.return_value = mock_session444        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')445        response = hook.describe_transform_job(job_name)446        mock_session.describe_transform_job.\447            assert_called_once_with(TransformJobName=job_name)448        self.assertEqual(response, 'InProgress')449    @mock.patch.object(SageMakerHook, 'get_conn')450    def test_describe_model(self, mock_client):451        mock_session = mock.Mock()452        attrs = {'describe_model.return_value':453                 model_name}454        mock_session.configure_mock(**attrs)455        mock_client.return_value = mock_session456        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')457        response = hook.describe_model(model_name)458        mock_session.describe_model.\459            assert_called_once_with(ModelName=model_name)460        self.assertEqual(response, model_name)461    @mock.patch.object(SageMakerHook, 'get_conn')462    def test_describe_endpoint_config(self, mock_client):463        mock_session = mock.Mock()464        attrs = {'describe_endpoint_config.return_value':465                 config_name}466        mock_session.configure_mock(**attrs)467        mock_client.return_value = mock_session468        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')469        response = hook.describe_endpoint_config(config_name)470        mock_session.describe_endpoint_config.\471            assert_called_once_with(EndpointConfigName=config_name)472        self.assertEqual(response, config_name)473    @mock.patch.object(SageMakerHook, 'get_conn')474    def test_describe_endpoint(self, mock_client):475        mock_session = mock.Mock()476        attrs = {'describe_endpoint.return_value':477                 'InProgress'}478        mock_session.configure_mock(**attrs)479        mock_client.return_value = mock_session480        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')481        response = hook.describe_endpoint(endpoint_name)482        mock_session.describe_endpoint.\483            assert_called_once_with(EndpointName=endpoint_name)484        self.assertEqual(response, 'InProgress')485    def test_secondary_training_status_changed_true(self):486        changed = secondary_training_status_changed(SECONDARY_STATUS_DESCRIPTION_1,487                                                    SECONDARY_STATUS_DESCRIPTION_2)488        self.assertTrue(changed)489    def test_secondary_training_status_changed_false(self):490        changed = secondary_training_status_changed(SECONDARY_STATUS_DESCRIPTION_1,491                                                    SECONDARY_STATUS_DESCRIPTION_1)492        self.assertFalse(changed)493    def test_secondary_training_status_message_status_changed(self):494        now = datetime.now(get_localzone())495        SECONDARY_STATUS_DESCRIPTION_1['LastModifiedTime'] = now496        expected = '{} {} - {}'.format(497            datetime.utcfromtimestamp(time.mktime(now.timetuple())).strftime('%Y-%m-%d %H:%M:%S'),498            status,499            message500        )501        self.assertEqual(502            secondary_training_status_message(SECONDARY_STATUS_DESCRIPTION_1, SECONDARY_STATUS_DESCRIPTION_2),503            expected)504    @mock.patch.object(SageMakerHook, 'get_log_conn')505    @mock.patch.object(SageMakerHook, 'get_conn')506    @mock.patch.object(time, 'time')507    def test_describe_training_job_with_logs_in_progress(self, mock_time, mock_client, mock_log_client):508        mock_session = mock.Mock()509        mock_log_session = mock.Mock()510        attrs = {'describe_training_job.return_value':511                 DESCRIBE_TRAINING_COMPELETED_RETURN512                 }513        log_attrs = {'describe_log_streams.side_effect':514                     LIFECYCLE_LOG_STREAMS,515                     'get_log_events.side_effect':516                     STREAM_LOG_EVENTS517                     }518        mock_time.return_value = 50519        mock_session.configure_mock(**attrs)520        mock_client.return_value = mock_session521        mock_log_session.configure_mock(**log_attrs)522        mock_log_client.return_value = mock_log_session523        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')524        response = hook.describe_training_job_with_log(job_name=job_name,525                                                       positions={},526                                                       stream_names=[],527                                                       instance_count=1,528                                                       state=LogState.WAIT_IN_PROGRESS,529                                                       last_description={},530                                                       last_describe_job_call=0)531        self.assertEqual(response, (LogState.JOB_COMPLETE, {}, 50))532    @mock.patch.object(SageMakerHook, 'get_log_conn')533    @mock.patch.object(SageMakerHook, 'get_conn')534    def test_describe_training_job_with_logs_job_complete(self, mock_client, mock_log_client):535        mock_session = mock.Mock()536        mock_log_session = mock.Mock()537        attrs = {'describe_training_job.return_value':538                 DESCRIBE_TRAINING_COMPELETED_RETURN539                 }540        log_attrs = {'describe_log_streams.side_effect':541                     LIFECYCLE_LOG_STREAMS,542                     'get_log_events.side_effect':543                     STREAM_LOG_EVENTS544                     }545        mock_session.configure_mock(**attrs)546        mock_client.return_value = mock_session547        mock_log_session.configure_mock(**log_attrs)548        mock_log_client.return_value = mock_log_session549        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')550        response = hook.describe_training_job_with_log(job_name=job_name,551                                                       positions={},552                                                       stream_names=[],553                                                       instance_count=1,554                                                       state=LogState.JOB_COMPLETE,555                                                       last_description={},556                                                       last_describe_job_call=0)557        self.assertEqual(response, (LogState.COMPLETE, {}, 0))558    @mock.patch.object(SageMakerHook, 'get_log_conn')559    @mock.patch.object(SageMakerHook, 'get_conn')560    def test_describe_training_job_with_logs_complete(self, mock_client, mock_log_client):561        mock_session = mock.Mock()562        mock_log_session = mock.Mock()563        attrs = {'describe_training_job.return_value':564                 DESCRIBE_TRAINING_COMPELETED_RETURN565                 }566        log_attrs = {'describe_log_streams.side_effect':567                     LIFECYCLE_LOG_STREAMS,568                     'get_log_events.side_effect':569                     STREAM_LOG_EVENTS570                     }571        mock_session.configure_mock(**attrs)572        mock_client.return_value = mock_session573        mock_log_session.configure_mock(**log_attrs)574        mock_log_client.return_value = mock_log_session575        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')576        response = hook.describe_training_job_with_log(job_name=job_name,577                                                       positions={},578                                                       stream_names=[],579                                                       instance_count=1,580                                                       state=LogState.COMPLETE,581                                                       last_description={},582                                                       last_describe_job_call=0)583        self.assertEqual(response, (LogState.COMPLETE, {}, 0))584    @mock.patch.object(SageMakerHook, 'check_training_config')585    @mock.patch.object(SageMakerHook, 'get_log_conn')586    @mock.patch.object(SageMakerHook, 'get_conn')587    @mock.patch.object(SageMakerHook, 'describe_training_job_with_log')588    def test_training_with_logs(self, mock_describe, mock_client, mock_log_client, mock_check_training):589        mock_check_training.return_value = True590        mock_describe.side_effect = \591            [(LogState.WAIT_IN_PROGRESS, DESCRIBE_TRAINING_INPROGRESS_RETURN, 0),592             (LogState.JOB_COMPLETE, DESCRIBE_TRAINING_STOPPING_RETURN, 0),593             (LogState.COMPLETE, DESCRIBE_TRAINING_COMPELETED_RETURN, 0)]594        mock_session = mock.Mock()595        mock_log_session = mock.Mock()596        attrs = {'create_training_job.return_value':597                 test_arn_return,598                 'describe_training_job.return_value':599                     DESCRIBE_TRAINING_COMPELETED_RETURN600                 }601        log_attrs = {'describe_log_streams.side_effect':602                     LIFECYCLE_LOG_STREAMS,603                     'get_log_events.side_effect':604                     STREAM_LOG_EVENTS605                     }606        mock_session.configure_mock(**attrs)607        mock_log_session.configure_mock(**log_attrs)608        mock_client.return_value = mock_session609        mock_log_client.return_value = mock_log_session610        hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id_1')611        hook.create_training_job(create_training_params, wait_for_completion=True,612                                 print_log=True, check_interval=1)613        self.assertEqual(mock_describe.call_count, 3)614        self.assertEqual(mock_session.describe_training_job.call_count, 1)615if __name__ == '__main__':...

test_mnm_drmaa_wrapper.py

Source:test_mnm_drmaa_wrapper.py

1from unittest.mock import Mock2import unittest3import drmaa4import mnm_drmaa_wrapper5from mnm_drmaa_wrapper import submit_drmaa_job, run_job_using_drmaa6# reduce waiting times for tests7mnm_drmaa_wrapper.GEVENT_TIMEOUT_WAIT = 08mnm_drmaa_wrapper.GEVENT_TIMEOUT_STARTUP = 09mnm_drmaa_wrapper.JOBSTATUS_FAILED_TIMEOUT = 010class SubmitDrmaaJobTests(unittest.TestCase):11    job_script_directory = "/tmp"12    exc_msg = "code 2: slurm_submit_batch_job error: Socket timed out on send/recv operation"13    def setUp(self) -> None:14        self.mock_logger = Mock()15        #self.mock_logger.debug.side_effect = lambda x : print("DEBUG:", x)16        #self.mock_logger.info.side_effect = lambda x: print("INFO:", x)17        self.mock_pipeline = Mock()18        self.mock_pipeline.is_job_suspended.return_value = False19        self.mock_session = Mock()20        self.mock_session.runJob.return_value=11111121    def test_no_exceptions_scenario(self) -> None:22        self.mock_session.jobStatus.side_effect = [drmaa.JobState.RUNNING, drmaa.JobState.DONE]23        jobid, _ = submit_drmaa_job("echo Hello", self.mock_session, None, self.mock_logger, self.mock_pipeline)24        self.mock_session.runJob.assert_called_once()25        self.assertEqual(self.mock_session.jobStatus.call_count, 2)26        self.assertEqual(jobid, self.mock_session.runJob.return_value)27    def test_job_status_2_exceptions(self) -> None:28        self.mock_session.jobStatus.side_effect = [ drmaa.errors.DrmCommunicationException(self.exc_msg),29                                                    drmaa.errors.DrmCommunicationException(self.exc_msg),30                                                    drmaa.JobState.RUNNING,31                                                    drmaa.JobState.DONE ]32        jobid, _ = submit_drmaa_job("echo Hello", self.mock_session, None, self.mock_logger, self.mock_pipeline)33        self.mock_session.runJob.assert_called_once()34        self.assertEqual(self.mock_session.jobStatus.call_count, 4)35        self.assertEqual(jobid, self.mock_session.runJob.return_value)36    def test_job_status_5_exceptions(self) -> None:37        self.mock_session.jobStatus.side_effect = [ drmaa.errors.DrmCommunicationException(self.exc_msg),38                                                    drmaa.errors.DrmCommunicationException(self.exc_msg),39                                                    drmaa.errors.DrmCommunicationException(self.exc_msg),40                                                    drmaa.errors.DrmCommunicationException(self.exc_msg),41                                                    drmaa.errors.DrmCommunicationException(self.exc_msg),42                                                    drmaa.JobState.RUNNING,43                                                    drmaa.JobState.DONE ]44        with self.assertRaises(drmaa.errors.DrmCommunicationException):45            submit_drmaa_job("echo Hello", self.mock_session, None, self.mock_logger, self.mock_pipeline)46        self.mock_session.runJob.assert_called_once()47        self.assertEqual(self.mock_session.jobStatus.call_count, 5)48class RunJobUsingDrmaaTests(unittest.TestCase):49    def setUp(self) -> None:50        self.mock_logger = Mock()51        #self.mock_logger.debug.side_effect = lambda x : print("DEBUG:", x)52        #self.mock_logger.info.side_effect = lambda x: print("INFO:", x)53        self.mock_pipeline = Mock()54        self.mock_pipeline.is_job_suspended.return_value = False55        self.mock_session = Mock()56        self.mock_session.jobStatus.side_effect = [drmaa.JobState.RUNNING,57                                                   drmaa.JobState.DONE]58        self.mock_jobid = 11111159        exc_msg = "code 2: slurm_submit_batch_job error: Socket timed out on send/recv operation"60        self.mock_session.runJob.side_effect = [ drmaa.errors.DrmCommunicationException(exc_msg),61                                                 drmaa.errors.DrmCommunicationException(exc_msg),62                                                 self.mock_jobid ]63        submit_drmaa_job = Mock()64    def test_run_job_using_drmaa_failing_to_submit_resubmit_0(self) -> None:65        with self.assertRaises(drmaa.errors.DrmCommunicationException):66            run_job_using_drmaa("echo Hello",67                                drmaa_session=self.mock_session,68                                logger = self.mock_logger,69                                resubmit = 0,70                                pipeline = self.mock_pipeline)71        self.mock_session.runJob.assert_called_once()72    def test_run_job_using_drmaa_failing_to_submit_resubmit_2(self) -> None:73        with self.assertRaises(mnm_drmaa_wrapper.error_drmaa_job):74            run_job_using_drmaa("echo Hello",75                                drmaa_session=self.mock_session,76                                logger=self.mock_logger,77                                resubmit=2,78                                pipeline=self.mock_pipeline)79        self.assertEqual(2, self.mock_session.runJob.call_count)80    def test_run_job_using_drmaa_failing_to_submit_resubmit_5(self) -> None:81        run_job_using_drmaa("echo Hello",82                            drmaa_session=self.mock_session,83                            logger=self.mock_logger,84                            resubmit=5,85                            pipeline=self.mock_pipeline)86        self.assertEqual(3, self.mock_session.runJob.call_count)87        self.assertEqual(2, self.mock_session.jobStatus.call_count)88if __name__ == '__main__':...

conftest.py

Source:conftest.py

1# coding: utf-82from __future__ import unicode_literals3import copy4import json5from mock import Mock, MagicMock6import pytest7from boxsdk.config import API, Client, Proxy8from boxsdk.network import default_network9from boxsdk.network.default_network import DefaultNetworkResponse, DefaultNetwork10from boxsdk.session.box_response import BoxResponse11from boxsdk.session.session import Session12from boxsdk.util.translator import Translator13@pytest.fixture(scope='function', autouse=True)14def original_default_translator():15    """A reference to the default translator, before the reference is changed by `default_translator` below."""16    return Translator._default_translator   # pylint:disable=protected-access17@pytest.yield_fixture(scope='function', autouse=True)18def default_translator(original_default_translator):19    """The default translator to use during the test.20    We don't want global state to mutate across tests. So before each test21    (because of autouse=True), we make a copy of the default translator, and22    assign this copy to Translator._default_translator. At the end of the test,23    we reset the reference.24    """25    try:26        translator = Translator(dict(copy.deepcopy(original_default_translator)), extend_default_translator=False, new_child=False)27        Translator._default_translator = translator   # pylint:disable=protected-access28        yield translator29    finally:30        Translator._default_translator = original_default_translator  # pylint:disable=protected-access31@pytest.fixture(scope='function')32def translator(default_translator):   # pylint:disable=unused-argument33    return Translator(extend_default_translator=True, new_child=True)34@pytest.fixture(scope='function')35def mock_box_session(translator):36    mock_session = MagicMock(Session)37    # pylint:disable=protected-access38    mock_session._api_config = mock_session.api_config = API()39    mock_session._client_config = mock_session.client_config = Client()40    mock_session._proxy_config = mock_session.proxy_config = Proxy()41    # pylint:enable=protected-access42    mock_session.get_url.side_effect = lambda *args, **kwargs: Session.get_url(mock_session, *args, **kwargs)43    mock_session.translator = translator44    return mock_session45@pytest.fixture()46def mock_box_session_2(translator):47    mock_session = MagicMock(Session)48    # pylint:disable=protected-access49    mock_session._api_config = API()50    mock_session._client_config = Client()51    mock_session._proxy_config = Proxy()52    # pylint:enable=protected-access53    mock_session.get_url.side_effect = lambda *args, **kwargs: Session.get_url(mock_session, *args, **kwargs)54    mock_session.translator = translator55    return mock_session56@pytest.fixture(scope='function')57def mock_network_layer():58    mock_network = Mock(DefaultNetwork)59    return mock_network60@pytest.fixture(autouse=True)61def prevent_tests_from_making_real_network_requests(monkeypatch):62    monkeypatch.delattr(default_network.requests.Session, 'request')63@pytest.fixture(scope='function')64def mock_user_response(mock_user_id, make_mock_box_request):65    # pylint:disable=redefined-outer-name66    mock_box_response, _ = make_mock_box_request(67        response={'type': 'user', 'id': mock_user_id},68    )69    return mock_box_response70@pytest.fixture(scope='function')71def mock_group_response(mock_group_id, make_mock_box_request):72    # pylint:disable=redefined-outer-name73    mock_box_response, _ = make_mock_box_request(74        response={'type': 'group', 'id': mock_group_id},75    )76    return mock_box_response77@pytest.fixture()78def make_mock_box_request():79    def inner(status_code=200, response_ok=True, response=None, content=None):80        mock_box_response = Mock(BoxResponse)81        mock_network_response = Mock(DefaultNetworkResponse)82        mock_box_response.network_response = mock_network_response83        mock_box_response.status_code = status_code84        mock_box_response.ok = response_ok85        if response is not None:86            mock_box_response.json.return_value = response87            mock_box_response.content = json.dumps(response).encode()88        else:89            mock_box_response.content = content90        return mock_box_response, mock_network_response91    return inner92@pytest.fixture(scope='function')93def mock_file_response(mock_object_id, make_mock_box_request):94    # pylint:disable=redefined-outer-name95    mock_box_response, _ = make_mock_box_request(96        response={'type': 'file', 'id': mock_object_id},97    )98    return mock_box_response99@pytest.fixture(scope='function')100def mock_folder_response(mock_object_id, make_mock_box_request):101    # pylint:disable=redefined-outer-name102    mock_box_response, _ = make_mock_box_request(103        response={'type': 'folder', 'id': mock_object_id},104    )105    return mock_box_response106@pytest.fixture(scope='function')107def mock_web_link_response(mock_object_id, make_mock_box_request):108    # pylint:disable=redefined-outer-name109    mock_box_response, _ = make_mock_box_request(110        response={'type': 'web_link', 'id': mock_object_id},111    )...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.