aws python boto3 s3 with assume_role

import boto3
from boto3.session import Session
 
role_arn = os.environ.get('ASSUME_ROLE')
 
if roleArn:
   client = boto3.client('sts')
   assumed_role = client.assume_role(RoleArn=roleArn, RoleSessionName='test')
   session = Session(aws_access_key_id=assumed_role['Credentials']['AccessKeyId'],
                      aws_secret_access_key=assumed_role['Credentials']['SecretAccessKey'],
                      aws_session_token=assumed_role['Credentials']['SessionToken'])
   s3_client = session.client('s3')
   s3 = session.resource('s3')
 
else:
   s3_client = boto3.client('s3')
   s3 = boto3.resource('s3')
 
s3.Bucket('myBucket').download_file('myKey', 'myFileName')
s3_client.upload_file('myKey', 'myBucket', 'myFileName')