From f3b0bccb0bea1840e9e566d19b4f249c9330fa39 Mon Sep 17 00:00:00 2001
From: Jeff Franklin <jpf@uw.edu>
Date: Tue, 5 Feb 2019 15:29:47 -0800
Subject: [PATCH] Get mock interface working better with Identity.UW.

---
 tests/test_mock.py |  5 +++--
 uw_saml2/mock.py   | 16 +++++++++++++---
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/tests/test_mock.py b/tests/test_mock.py
index 8e1501c..6527ee4 100644
--- a/tests/test_mock.py
+++ b/tests/test_mock.py
@@ -8,13 +8,14 @@ def flick_mock_flag(monkeypatch):
 
 
 def test_login_redirect():
-    expected = ('mock-login?return_to=%2F&force_authn=False'
+    expected = ('/mock-login?return_to=%2F&force_authn=False'
                 '&idp=urn%3Amace%3Aincommon%3Awashington.edu')
     assert uw_saml2.login_redirect() == expected
 
 
 def test_process_response():
-    post = {'idp': 'urn:mace:incommon:washington.edu', 'foo': 'bar'}
+    post = {'idp': 'urn:mace:incommon:washington.edu',
+            'remote_user': 'javerage@washington.edu', 'uwnetid': 'javerage'}
     expected = {'two_factor': False, **post}
     assert uw_saml2.process_response(post) == expected
 
diff --git a/uw_saml2/mock.py b/uw_saml2/mock.py
index c6cde6c..8fc5bba 100644
--- a/uw_saml2/mock.py
+++ b/uw_saml2/mock.py
@@ -1,4 +1,6 @@
 import urllib.parse
+from .idp import uw, federated
+MOCK_LOGIN_URL = '/mock-login'
 
 
 class SamlAuthenticator:
@@ -10,7 +12,7 @@ def __init__(self, request, old_settings, *args, **kwargs):
     def login(self, **kwargs):
         kwargs['idp'] = self.idp
         qs = urllib.parse.urlencode(kwargs)
-        return f'mock-login?{qs}'
+        return f'{MOCK_LOGIN_URL}?{qs}'
 
     def process_response(self):
         idp = self.request['post_data'].get('idp')
@@ -20,8 +22,16 @@ def process_response(self):
 
     def get_attributes(self):
         """Just reflect what's posted right back."""
-        items = self.request['post_data'].items()
-        return dict((key, [value]) for key, value in items)
+        post = self.request['post_data']
+        remote_user = post.get('remote_user', '')
+        attributes = {}
+        if post.get('idp') == uw.UwIdp.entity_id:
+            uwnetid = remote_user.split('@washington.edu')[0]
+            attributes['uwnetid'] = [uwnetid]
+        elif post.get('idp') == federated.CollegenetIdp.entity_id:
+            attributes['collegenet_userid'] = [remote_user]
+        attributes.update((key, [value]) for key, value in post.items())
+        return attributes
 
     def get_errors(self):
         return self.errors