Update tests

tests/test_attributes.py

@@ -48,56 +48,69 @@ def _make_test_saml_response(
             }
 
             if display_name:
-                self.ava["displayName"] = display_name
+                self.ava["displayName"] = [display_name]
 
     return FakeResponse()
 
 
-class SamlUsernameTestCase(unittest.TestCase):
+class SamlUserAttributeTestCase(unittest.TestCase):
 
-    def test_normal_user(self):
+    def _attribute_test(
+        self,
+        input_uid: str,
+        input_displayname: Optional[str],
+        output_localpart: str,
+        output_displayname: Optional[str],
+    ):
+        """Creates a dummy response, feeds it to the provider and checks the output
+
+        Args:
+            input_uid: The value of the mxid_source_attribute that the provider will
+                base the generated localpart off of.
+            input_displayname: The saml auth response displayName value that the
+                provider will generate a Matrix user displayname from.
+            output_localpart: The expected mxid localpart.
+            output_displayname: The expected matrix displayname.
+
+        """
         provider, config = create_mapping_provider()
-        response = _make_test_saml_response(config, "john*doe2000#@example.com", None)
+        response = _make_test_saml_response(config, input_uid, input_displayname)
 
         attribute_dict = provider.saml_response_to_user_attributes(response)
-        self.assertEqual(attribute_dict["mxid_localpart"], "john.doe2000")
-        self.assertEqual(attribute_dict["displayname"], "john.doe2000")
+        self.assertEqual(attribute_dict["mxid_localpart"], output_localpart)
+        self.assertEqual(attribute_dict["displayname"], output_displayname)
 
-    def test_multiple_adjacent_symbols(self):
-        provider = create_mapping_provider()
+    def test_normal_user(self):
+        self._attribute_test("john*doe2000#@example.com", None, "john.doe2000", None)
 
-        username = "bob%^$&#!bobby@example.com"
-        localpart = provider.saml_response_to_user_attributes(username)
-        self.assertEqual(localpart, "bob.bobby")
+    def test_normal_user_displayname(self):
+        self._attribute_test(
+            "john*doe2000#@example.com", "Jonny", "john.doe2000", "Jonny"
+        )
+
+    def test_multiple_adjacent_symbols(self):
+        self._attribute_test("bob%^$&#!bobby@example.com", None, "bob.bobby", None)
 
     def test_username_does_not_end_with_dot(self):
         """This is allowed in mxid syntax, but is not aesthetically pleasing"""
-        provider = create_mapping_provider()
-
-        username = "bob.bobby$@example.com"
-        localpart = provider.saml_response_to_user_attributes(username)
-        self.assertEqual(localpart, "bob.bobby")
+        self._attribute_test("bob.bobby$@example.com", None, "bob.bobby", None)
 
     def test_username_no_email(self):
-        provider = create_mapping_provider()
-
-        username = "bob.bobby"
-        localpart = provider.saml_response_to_user_attributes(username)
-        self.assertEqual(localpart, "bob.bobby")
+        self._attribute_test("bob.bobby", None, "bob.bobby", None)
 
     def test_username_starting_with_underscore(self):
-        provider = create_mapping_provider()
-
-        username = "_twilight (sparkle)@somewhere.com"
-        localpart = provider.saml_response_to_user_attributes(username)
-        self.assertEqual(localpart, "twilight.sparkle")
+        self._attribute_test(
+            "_twilight (sparkle)@somewhere.com", None, "twilight.sparkle", None
+        )
 
     def test_existing_user(self):
-        provider = create_mapping_provider()
+        provider, config = create_mapping_provider()
+        response = _make_test_saml_response(config, "wibble%@wobble.com", None)
 
-        username = "wibble%@wobble.com"
-        localpart = provider.saml_response_to_user_attributes(username)
+        attribute_dict = provider.saml_response_to_user_attributes(response)
 
         # Simulate a failure on the first attempt
-        localpart = provider.saml_response_to_user_attributes(username, failures=1)
-        self.assertEqual(localpart, "wibble1")
+        attribute_dict = provider.saml_response_to_user_attributes(response, failures=1)
+
+        self.assertEqual(attribute_dict["mxid_localpart"], "wibble1")
+        self.assertEqual(attribute_dict["displayname"], None)