Improve UsernamePickerTestCase

tests/handlers/test_oidc.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
-import re
 from typing import Dict, Optional
 from urllib.parse import parse_qs, urlencode, urlparse
 
@@ -23,15 +22,15 @@
 
 from twisted.web.resource import Resource
 
-from synapse.api.errors import RedirectException
 from synapse.handlers.sso import MappingException
 from synapse.rest.client.v1 import login
 from synapse.rest.synapse.client.pick_username import pick_username_resource
 from synapse.server import HomeServer
 from synapse.types import UserID
 
+from tests.rest.client.v1.utils import TEST_OIDC_CONFIG
 from tests.test_utils import FakeResponse, simple_async_mock
-from tests.unittest import HomeserverTestCase, override_config
+from tests.unittest import HomeserverTestCase, override_config, skip_unless
 
 try:
     import authlib  # noqa: F401
@@ -855,69 +854,60 @@ def _generate_oidc_session_token(
         )
 
 
+@skip_unless(HAS_OIDC, "requires OIDC")
 class UsernamePickerTestCase(HomeserverTestCase):
-    if not HAS_OIDC:
-        skip = "requires OIDC"
+    """Tests for the username picker flow of SSO login"""
 
     servlets = [login.register_servlets]
 
     def default_config(self):
         config = super().default_config()
         config["public_baseurl"] = BASE_URL
-        oidc_config = {
-            "enabled": True,
-            "client_id": CLIENT_ID,
-            "client_secret": CLIENT_SECRET,
-            "issuer": ISSUER,
-            "scopes": SCOPES,
-            "user_mapping_provider": {
-                "config": {"display_name_template": "{{ user.displayname }}"}
-            },
-        }
 
-        # Update this config with what's in the default config so that
-        # override_config works as expected.
-        oidc_config.update(config.get("oidc_config", {}))
-        config["oidc_config"] = oidc_config
+        config["oidc_config"] = {}
+        config["oidc_config"].update(TEST_OIDC_CONFIG)
+        config["oidc_config"]["user_mapping_provider"] = {
+            "config": {"display_name_template": "{{ user.displayname }}"}
+        }
 
         # whitelist this client URI so we redirect straight to it rather than
         # serving a confirmation page
         config["sso"] = {"client_whitelist": ["https://whitelisted.client"]}
         return config
 
     def create_resource_dict(self) -> Dict[str, Resource]:
+        from synapse.rest.oidc import OIDCResource
+
         d = super().create_resource_dict()
         d["/_synapse/client/pick_username"] = pick_username_resource(self.hs)
+        d["/_synapse/oidc"] = OIDCResource(self.hs)
         return d
 
     def test_username_picker(self):
         """Test the happy path of a username picker flow."""
         client_redirect_url = "https://whitelisted.client"
 
-        # first of all, mock up an OIDC callback to the OidcHandler, which should
-        # raise a RedirectException
-        userinfo = {"sub": "tester", "displayname": "Jonny"}
-        f = self.get_failure(
-            _make_callback_with_userinfo(
-                self.hs, userinfo, client_redirect_url=client_redirect_url
-            ),
-            RedirectException,
+        # do the start of the login flow
+        channel = self.helper.auth_via_oidc(
+            {"sub": "tester", "displayname": "Jonny"}, client_redirect_url
         )
 
-        # check the Location and cookies returned by the RedirectException
-        self.assertEqual(f.value.location, b"/_synapse/client/pick_username")
-        cookieheader = f.value.cookies[0]
-        regex = re.compile(b"^username_mapping_session=([a-zA-Z]+);")
-        m = regex.search(cookieheader)
-        if not m:
-            self.fail("cookie header %s does not match %s" % (cookieheader, regex))
+        # that should redirect to the username picker
+        self.assertEqual(channel.code, 302, channel.result)
+        picker_url = channel.headers.getRawHeaders("Location")[0]
+        self.assertEqual(picker_url, "/_synapse/client/pick_username")
+
+        # ... with a username_mapping_session cookie
+        cookies = {}  # type: Dict[str,str]
+        channel.extract_cookies(cookies)
+        self.assertIn("username_mapping_session", cookies)
+        session_id = cookies["username_mapping_session"]
 
         # introspect the sso handler a bit to check that the username mapping session
         # looks ok.
-        session_id = m.group(1).decode("ascii")
         username_mapping_sessions = self.hs.get_sso_handler()._username_mapping_sessions
         self.assertIn(
-            session_id, username_mapping_sessions, "session id not found in map"
+            session_id, username_mapping_sessions, "session id not found in map",
         )
         session = username_mapping_sessions[session_id]
         self.assertEqual(session.remote_user_id, "tester")
@@ -930,15 +920,15 @@ def test_username_picker(self):
 
         # Now, submit a username to the username picker, which should serve a redirect
         # back to the client
-        submit_path = f.value.location + b"/submit"
+        submit_path = picker_url + "/submit"
         content = urlencode({b"username": b"bobby"}).encode("utf8")
         chan = self.make_request(
             "POST",
             path=submit_path,
             content=content,
             content_is_form=True,
             custom_headers=[
-                ("Cookie", cookieheader),
+                ("Cookie", "username_mapping_session=" + session_id),
                 # old versions of twisted don't do form-parsing without a valid
                 # content-length header.
                 ("Content-Length", str(len(content))),

tests/rest/client/v1/utils.py

@@ -363,7 +363,7 @@ def login_via_oidc(self, remote_user_id: str) -> JsonDict:
         the normal places.
         """
         client_redirect_url = "https://x"
-        channel = self.auth_via_oidc(remote_user_id, client_redirect_url)
+        channel = self.auth_via_oidc({"sub": remote_user_id}, client_redirect_url)
 
         # expect a confirmation page
         assert channel.code == 200
@@ -390,7 +390,7 @@ def login_via_oidc(self, remote_user_id: str) -> JsonDict:
 
     def auth_via_oidc(
         self,
-        remote_user_id: str,
+        user_info_dict: JsonDict,
         client_redirect_url: Optional[str] = None,
         ui_auth_session_id: Optional[str] = None,
     ) -> FakeChannel:
@@ -411,7 +411,8 @@ def auth_via_oidc(
         the normal places.
 
         Args:
-            remote_user_id: the remote id that the OIDC provider should present
+            user_info_dict: the remote userinfo that the OIDC provider should present.
+                Typically this should be '{"sub": "<remote user id>"}'.
             client_redirect_url: for a login flow, the client redirect URL to pass to
                 the login redirect endpoint
             ui_auth_session_id: if set, we will perform a UI Auth flow. The session id
@@ -457,7 +458,7 @@ def auth_via_oidc(
             # a dummy OIDC access token
             ("https://issuer.test/token", {"access_token": "TEST"}),
             # and then one to the user_info endpoint, which returns our remote user id.
-            ("https://issuer.test/userinfo", {"sub": remote_user_id}),
+            ("https://issuer.test/userinfo", user_info_dict),
         ]
 
         async def mock_req(method: str, uri: str, data=None, headers=None):

tests/rest/client/v2_alpha/test_auth.py

@@ -411,7 +411,7 @@ def test_ui_auth_via_sso(self):
         # run the UIA-via-SSO flow
         session_id = channel.json_body["session"]
         channel = self.helper.auth_via_oidc(
-            remote_user_id=remote_user_id, ui_auth_session_id=session_id
+            {"sub": remote_user_id}, ui_auth_session_id=session_id
         )
 
         # that should serve a confirmation page