Add test for websocket aws-cn fix

Author: jamesls

chalice/app.py

@@ -611,9 +611,13 @@ def configure_from_api_id(self, api_id, stage):
         region_name = self._get_region()
 
         if region_name.startswith("cn-"):
-            domain_name_template = '{api_id}.execute-api.{region}.amazonaws.com.cn'
+            domain_name_template = (
+                '{api_id}.execute-api.{region}.amazonaws.com.cn'
+            )
         else:
-            domain_name_template = '{api_id}.execute-api.{region}.amazonaws.com'
+            domain_name_template = (
+                '{api_id}.execute-api.{region}.amazonaws.com'
+            )
 
         domain_name = domain_name_template.format(
             api_id=api_id, region=region_name)

tests/unit/conftest.py

@@ -183,11 +183,13 @@ def create_event_inner(uri, method, path, content_type='application/json'):
 
 @fixture
 def create_websocket_event():
-    def create_event_inner(route_key, body=''):
+    def create_event_inner(
+            route_key, body='',
+            endpoint='abcd1234.execute-api.us-west-2.amazonaws.com'):
         return {
             'requestContext': {
                 'routeKey': route_key,
-                'domainName': 'abcd1234.execute-api.us-west-2.amazonaws.com',
+                'domainName': endpoint,
                 'stage': 'api',
                 'connectionId': 'ABCD1234=',
                 'apiId': 'abcd1234',

tests/unit/test_app.py

@@ -2824,6 +2824,28 @@ def test_can_configure_client_on_connect(sample_websocket_app,
     ]
 
 
+def test_can_configure_non_aws_partition_clients(sample_websocket_app,
+                                                 create_websocket_event,
+                                                 monkeypatch):
+    # tests/conftest.py already monkeypatches out environment variables,
+    # so if we want a test with a specific region we have to use the
+    # same approach and monkeypatch the os.environ.
+    monkeypatch.setenv('AWS_REGION', 'cn-north-1')
+    demo, _ = sample_websocket_app
+    client = FakeClient()
+    demo.websocket_api.session = FakeSession(client)
+    event = create_websocket_event(
+        '$connect',
+        endpoint='abcd1234.execute-api.cn-north-1.amazonaws.com.cn')
+    handler = websocket_handler_for_route('$connect', demo)
+    handler(event, context=None)
+
+    assert demo.websocket_api.session.calls == [
+        ('apigatewaymanagementapi',
+         'https://abcd1234.execute-api.cn-north-1.amazonaws.com.cn/api'),
+    ]
+
+
 def test_uses_api_id_not_domain_name(sample_websocket_app,
                                      create_websocket_event):
     demo, calls = sample_websocket_app