# Buggy code
def update_domain_name(
    self,
    protocol,  # type: str
    domain_name,  # type: str
    endpoint_type,  # type: str
    certificate_arn,  # type: str
    security_policy=None,  # type: Optional[str]
    tags=None,  # type: StrMap
):
    # type: (...) -> Dict[str, Any]
    if protocol == "HTTP":
        patch_operations = self.get_custom_domain_patch_operations(
            certificate_arn,
            endpoint_type,
            security_policy,
        )
        updated_domain_name = self._update_domain_name(domain_name, patch_operations)
    elif protocol == "WEBSOCKET":
        kwargs = self.get_custom_domain_params_v2(
            domain_name=domain_name,
            endpoint_type=endpoint_type,
            security_policy=security_policy,
            certificate_arn=certificate_arn,
        )
        updated_domain_name = self._update_domain_name_v2(kwargs)
    else:
        raise ValueError("Unsupported protocol value.")
    resource_arn = (
        "arn:aws:apigateway:{region_name}::/domainnames/{domain_name}".format(
            region_name=self.region_name, domain_name=domain_name
        )
    )
    self._update_resource_tags(resource_arn, tags)
    return updated_domain_name


# Traceback (KeyError)
Updating custom domain name: api.transferbank.com.br
Traceback (most recent call last):
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/chalice/cli/__init__.py", line 646, in main
return cli(obj={})
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/click/core.py", line 829, in __call__
return self.main(*args, **kwargs)
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/click/core.py", line 782, in main
rv = self.invoke(ctx)
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/click/core.py", line 1259, in invoke
return _process_result(sub_ctx.command.invoke(sub_ctx))
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/click/core.py", line 1066, in invoke
return ctx.invoke(self.callback, **ctx.params)
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/click/core.py", line 610, in invoke
return callback(*args, **kwargs)
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/click/decorators.py", line 21, in new_func
return f(get_current_context(), *args, **kwargs)
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/chalice/cli/__init__.py", line 206, in deploy
deployed_values = d.deploy(config, chalice_stage_name=stage)
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/chalice/deploy/deployer.py", line 374, in deploy
return self._deploy(config, chalice_stage_name)
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/chalice/deploy/deployer.py", line 390, in _deploy
self._executor.execute(plan)
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/chalice/deploy/executor.py", line 42, in execute
getattr(self, '_do_%s' % instruction.__class__.__name__.lower(),
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/chalice/deploy/executor.py", line 55, in _do_apicall
result = method(**final_kwargs)
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/chalice/awsclient.py", line 702, in update_domain_name
updated_domain_name = self._update_domain_name_v2(kwargs)
File "/Users/rafagan/.virtualenvs/tb-websocket-ws/lib/python3.8/site-packages/chalice/awsclient.py", line 802, in _update_domain_name_v2
'hosted_zone_id': result_data['HostedZoneId'],
KeyError: 'HostedZoneId'

# Fixed code
def update_domain_name(
    self,
    protocol,  # type: str
    domain_name,  # type: str
    endpoint_type,  # type: str
    certificate_arn,  # type: str
    security_policy=None,  # type: Optional[str]
    tags=None,  # type: StrMap
):
    # type: (...) -> DomainNameResponse
    if protocol == "HTTP":
        patch_operations = self.get_custom_domain_patch_operations(
            certificate_arn,
            endpoint_type,
            security_policy,
        )
        updated_domain_name = self._update_domain_name(domain_name, patch_operations)
    elif protocol == "WEBSOCKET":
        kwargs = self.get_custom_domain_params_v2(
            domain_name=domain_name,
            endpoint_type=endpoint_type,
            security_policy=security_policy,
            certificate_arn=certificate_arn,
        )
        updated_domain_name = self._update_domain_name_v2(kwargs)
    else:
        raise ValueError("Unsupported protocol value.")
    resource_arn = (
        "arn:aws:apigateway:{region_name}::/domainnames/{domain_name}".format(
            region_name=self.region_name, domain_name=domain_name
        )
    )
    self._update_resource_tags(resource_arn, tags)
    return updated_domain_name