# Buggy code
def create_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":
        kwargs = {
            "domainName": domain_name,
            "endpointConfiguration": {
                "types": [endpoint_type],
            },
        }
        if security_policy is not None:
            kwargs["securityPolicy"] = security_policy
        if endpoint_type == "EDGE":
            kwargs["certificateArn"] = certificate_arn
        else:
            kwargs["regionalCertificateArn"] = certificate_arn
        if tags is not None:
            kwargs["tags"] = tags
        created_domain_name = self._create_domain_name(kwargs)
    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,
            tags=tags,
        )
        created_domain_name = self._create_domain_name_v2(kwargs)
    else:
        raise ValueError("Unsupported protocol value.")
    return created_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 create_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":
        kwargs = {
            "domainName": domain_name,
            "endpointConfiguration": {
                "types": [endpoint_type],
            },
        }
        if security_policy is not None:
            kwargs["securityPolicy"] = security_policy
        if endpoint_type == "EDGE":
            kwargs["certificateArn"] = certificate_arn
        else:
            kwargs["regionalCertificateArn"] = certificate_arn
        if tags is not None:
            kwargs["tags"] = tags
        created_domain_name = self._create_domain_name(kwargs)
    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,
            tags=tags,
        )
        created_domain_name = self._create_domain_name_v2(kwargs)
    else:
        raise ValueError("Unsupported protocol value.")
    return created_domain_name