Repository: aws/chalice
Issue: Error creating custom domain for websockets

When I try to create a custom domain for my websocket the following error occurs:

```
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'
```

My configuration is:

```
"websocket_api_custom_domain": {
  "domain_name": "api.transferbank.com.br",
  "certificate_arn": "arn:aws:acm:us-east-1:073289305372:certificate/d995959e-6f7d-4fcb-9c56-bf370ca1a382",
  "url_prefix": "tb_wst_channel_manager"
},
```

Changed files:
tests/functional/test_awsclient.py, chalice/awsclient.py

diff --git a/chalice/awsclient.py b/chalice/awsclient.py
index 79c66ed..23737a4 100644
--- a/chalice/awsclient.py
+++ b/chalice/awsclient.py
@@ -63,6 +63,15 @@ LogEventsResponse = TypedDict(
         'nextToken': str,
     }, total=False
 )
+DomainNameResponse = TypedDict(
+    'DomainNameResponse', {
+        'domain_name': str,
+        'security_policy': str,
+        'hosted_zone_id': str,
+        'certificate_arn': str,
+        'alias_domain_name': str,
+    }
+)
 
 _REMOTE_CALL_ERRORS = (
     botocore.exceptions.ClientError, RequestsConnectionError
@@ -462,7 +471,7 @@ class TypedAWSClient(object):
                            security_policy=None,  # type: Optional[str]
                            tags=None,             # type: StrMap
                            ):
-        # type: (...) -> Dict[str, Any]
+        # type: (...) -> DomainNameResponse
         if protocol == 'HTTP':
             kwargs = {
                 'domainName': domain_name,
@@ -493,7 +502,7 @@ class TypedAWSClient(object):
         return created_domain_name
 
     def _create_domain_name(self, api_args):
-        # type: (Dict[str, Any]) -> Dict[str, Any]
+        # type: (Dict[str, Any]) -> DomainNameResponse
         client = self._client('apigateway')
         exceptions = (
             client.exceptions.TooManyRequestsException,
@@ -504,29 +513,31 @@ class TypedAWSClient(object):
             should_retry=lambda x: True,
             retryable_exceptions=exceptions
         )
-        domain_name = {
-            'domain_name': result['domainName'],
-            'endpoint_configuration': result['endpointConfiguration'],
-            'security_policy': result['securityPolicy'],
-        }
         if result.get('regionalHostedZoneId'):
-            domain_name['hosted_zone_id'] = result['regionalHostedZoneId']
+            hosted_zone_id = result['regionalHostedZoneId']
         else:
-            domain_name['hosted_zone_id'] = result['distributionHostedZoneId']
+            hosted_zone_id = result['distributionHostedZoneId']
 
         if result.get('regionalCertificateArn'):
-            domain_name['certificate_arn'] = result['regionalCertificateArn']
+            certificate_arn = result['regionalCertificateArn']
         else:
-            domain_name['certificate_arn'] = result['certificateArn']
+            certificate_arn = result['certificateArn']
 
         if result.get('regionalDomainName') is not None:
-            domain_name['alias_domain_name'] = result['regionalDomainName']
+            alias_domain_name = result['regionalDomainName']
         else:
-            domain_name['alias_domain_name'] = result['distributionDomainName']
+            alias_domain_name = result['distributionDomainName']
+        domain_name = {
+            'domain_name': result['domainName'],
+            'security_policy': result['securityPolicy'],
+            'hosted_zone_id': hosted_zone_id,
+            'certificate_arn': certificate_arn,
+            'alias_domain_name': alias_domain_name,
+        }  # type: DomainNameResponse
         return domain_name
 
     def _create_domain_name_v2(self, api_args):
-        # type: (Dict[str, Any]) -> Dict[str, Any]
+        # type: (Dict[str, Any]) -> DomainNameResponse
         client = self._client('apigatewayv2')
         exceptions = (
             client.exceptions.TooManyRequestsException,
@@ -539,12 +550,12 @@ class TypedAWSClient(object):
         )
         result_data = result['DomainNameConfigurations'][0]
         domain_name = {
-            'domain_name': result_data['ApiGatewayDomainName'],
-            'endpoint_type': result_data['EndpointType'],
+            'domain_name': result['DomainName'],
+            'alias_domain_name': result_data['ApiGatewayDomainName'],
             'security_policy': result_data['SecurityPolicy'],
             'hosted_zone_id': result_data['HostedZoneId'],
             'certificate_arn': result_data['CertificateArn']
-        }
+        }  # type: DomainNameResponse
         return domain_name
 
     def _create_lambda_function(self, api_args):
@@ -686,7 +697,7 @@ class TypedAWSClient(object):
                            security_policy=None,      # type: Optional[str]
                            tags=None,                 # type: StrMap
                            ):
-        # type: (...) -> Dict[str, Any]
+        # type: (...) -> DomainNameResponse
         if protocol == 'HTTP':
             patch_operations = self.get_custom_domain_patch_operations(
                 certificate_arn,
@@ -745,7 +756,7 @@ class TypedAWSClient(object):
                 ResourceArn=resource_arn, Tags=tags_to_add)
 
     def _update_domain_name(self, custom_domain_name, patch_operations):
-        # type: (str, List[Dict[str, str]]) -> Dict[str, Any]
+        # type: (str, List[Dict[str, str]]) -> DomainNameResponse
         client = self._client('apigateway')
         exceptions = (
             client.exceptions.TooManyRequestsException,
@@ -764,29 +775,31 @@ class TypedAWSClient(object):
             )
             result.update(response)
 
-        domain_name = {
-            'domain_name': result['domainName'],
-            'endpoint_configuration': result['endpointConfiguration'],
-            'security_policy': result['securityPolicy'],
-        }
         if result.get('regionalCertificateArn'):
-            domain_name['certificate_arn'] = result['regionalCertificateArn']
+            certificate_arn = result['regionalCertificateArn']
         else:
-            domain_name['certificate_arn'] = result['certificateArn']
+            certificate_arn = result['certificateArn']
 
         if result.get('regionalHostedZoneId'):
-            domain_name['hosted_zone_id'] = result['regionalHostedZoneId']
+            hosted_zone_id = result['regionalHostedZoneId']
         else:
-            domain_name['hosted_zone_id'] = result['distributionHostedZoneId']
+            hosted_zone_id = result['distributionHostedZoneId']
 
         if result.get('regionalDomainName') is not None:
-            domain_name['alias_domain_name'] = result['regionalDomainName']
+            alias_domain_name = result['regionalDomainName']
         else:
-            domain_name['alias_domain_name'] = result['distributionDomainName']
+            alias_domain_name = result['distributionDomainName']
+        domain_name = {
+            'domain_name': result['domainName'],
+            'security_policy': result['securityPolicy'],
+            'certificate_arn': certificate_arn,
+            'hosted_zone_id': hosted_zone_id,
+            'alias_domain_name': alias_domain_name
+        }  # type: DomainNameResponse
         return domain_name
 
     def _update_domain_name_v2(self, api_args):
-        # type: (Dict[str, Any]) -> Dict[str, Any]
+        # type: (Dict[str, Any]) -> DomainNameResponse
         client = self._client('apigatewayv2')
         exceptions = (
             client.exceptions.TooManyRequestsException,
@@ -801,11 +814,11 @@ class TypedAWSClient(object):
         result_data = result['DomainNameConfigurations'][0]
         domain_name = {
             'domain_name': result['DomainName'],
-            'endpoint_configuration': result_data['EndpointType'],
+            'alias_domain_name': result_data['ApiGatewayDomainName'],
             'security_policy': result_data['SecurityPolicy'],
             'hosted_zone_id': result_data['HostedZoneId'],
             'certificate_arn': result_data['CertificateArn']
-        }
+        }  # type: DomainNameResponse
         return domain_name
 
     def delete_domain_name(self, domain_name):
diff --git a/tests/functional/test_awsclient.py b/tests/functional/test_awsclient.py
index 5b9b35b..97d5fc9 100644
--- a/tests/functional/test_awsclient.py
+++ b/tests/functional/test_awsclient.py
@@ -778,9 +778,6 @@ class TestCreateDomainName(object):
             }
         ) == {
             'domain_name': 'test_domain',
-            'endpoint_configuration': {
-                'types': ['REGIONAL'],
-            },
             'security_policy': 'TLS_1_0',
             'alias_domain_name': 'regional_domain_name',
             'hosted_zone_id': 'hosted_zone_id',
@@ -833,9 +830,6 @@ class TestCreateDomainName(object):
             }
         ) == {
             'domain_name': 'test_domain',
-            'endpoint_configuration': {
-                'types': ['EDGE'],
-            },
             'security_policy': 'TLS_1_0',
             'hosted_zone_id': 'hosted_zone_id',
             'alias_domain_name': 'dist_test_domain',
@@ -862,7 +856,7 @@ class TestCreateDomainName(object):
                 'DomainName': 'test_websocket_domain',
                 'DomainNameConfigurations': [
                     {
-                        'ApiGatewayDomainName': 'test_websocket_domain',
+                        'ApiGatewayDomainName': 'd-1234',
                         'CertificateArn': 'certificate_arn',
                         'CertificateName': 'certificate_name',
                         'CertificateUploadDate': datetime.datetime.now(),
@@ -891,7 +885,7 @@ class TestCreateDomainName(object):
             }
         ) == {
             'domain_name': 'test_websocket_domain',
-            'endpoint_type': 'REGIONAL',
+            'alias_domain_name': 'd-1234',
             'security_policy': 'TLS_1_2',
             'hosted_zone_id': 'hosted_zone_id',
             'certificate_arn': 'certificate_arn'
@@ -1013,7 +1007,7 @@ class TestUpdateDomainName(object):
             'DomainName': 'test_domain',
             'DomainNameConfigurations': [
                 {
-                    'ApiGatewayDomainName': 'test_domain',
+                    'ApiGatewayDomainName': 'd-1234',
                     'CertificateArn': 'certificate_arn',
                     'CertificateName': 'certificate_name',
                     'CertificateUploadDate': datetime.datetime.now(),
@@ -1039,11 +1033,11 @@ class TestUpdateDomainName(object):
             security_policy='TLS_1_2',
             certificate_arn='certificate_arn',
         ) == {
-           'domain_name': 'test_domain',
-           'endpoint_configuration': 'REGIONAL',
-           'security_policy': 'TLS_1_2',
-           'hosted_zone_id': 'hosted_zone_id',
-           'certificate_arn': 'certificate_arn'
+            'domain_name': 'test_domain',
+            'alias_domain_name': 'd-1234',
+            'security_policy': 'TLS_1_2',
+            'hosted_zone_id': 'hosted_zone_id',
+            'certificate_arn': 'certificate_arn'
         }
         stubbed_session.verify_stubs()
 
@@ -1229,11 +1223,6 @@ class TestUpdateDomainName(object):
             certificate_arn='regional_certificate_arn',
         ) == {
             'domain_name': 'test_domain',
-            'endpoint_configuration': {
-                'types': [
-                    'REGIONAL',
-                ],
-            },
             'alias_domain_name': 'regional_domain_name',
             'security_policy': 'TLS_1_2',
             'hosted_zone_id': 'hosted_zone_id',
@@ -1307,11 +1296,6 @@ class TestUpdateDomainName(object):
             certificate_arn='certificate_arn',
         ) == {
             'domain_name': 'test_domain',
-            'endpoint_configuration': {
-                 'types': [
-                     'EDGE',
-                 ],
-             },
             'security_policy': 'TLS_1_0',
             'alias_domain_name': 'dist_domain_name',
             'hosted_zone_id': 'hosted_zone_id',