KaneAI - World’s First E2E Software Testing Agent.Book a Demo 
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.
LambdaTest Learning Hubs:
- JUnit Tutorial
- TestNG Tutorial
- Webdriver Tutorial
- WebDriverIO Tutorial
- Protractor Tutorial
- Selenium 4 Tutorial
- Jenkins Tutorial
- NUnit Tutorial
- Jest Tutorial
- Playwright Tutorial
- Cypress Tutorial
- PyTest Tutorial
YouTube
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Run tavern automation tests on LambdaTest cloud grid
Perform automation testing on 3000+ real desktop and mobile devices online.
Try LambdaTest Now !!
Get 100 minutes of automation test minutes FREE!!
Was this article helpful?
