Update the upstream oauth user registration test

Author: sandhose

crates/handlers/src/upstream_oauth2/link.rs

@@ -1075,20 +1075,18 @@ mod tests {
     use mas_data_model::{
         UpstreamOAuthAuthorizationSession, UpstreamOAuthLink, UpstreamOAuthProviderClaimsImports,
         UpstreamOAuthProviderImportPreference, UpstreamOAuthProviderLocalpartPreference,
-        UpstreamOAuthProviderTokenAuthMethod,
+        UpstreamOAuthProviderTokenAuthMethod, UserEmailAuthentication, UserRegistration,
     };
     use mas_iana::jose::JsonWebSignatureAlg;
     use mas_jose::jwt::{JsonWebSignatureHeader, Jwt};
     use mas_keystore::Keystore;
     use mas_router::Route;
-    use mas_storage::{
-        Pagination, Repository, RepositoryError, upstream_oauth2::UpstreamOAuthProviderParams,
-        user::UserEmailFilter,
-    };
+    use mas_storage::{Repository, RepositoryError, upstream_oauth2::UpstreamOAuthProviderParams};
     use oauth2_types::scope::{OPENID, Scope};
     use rand_chacha::ChaChaRng;
     use serde_json::Value;
     use sqlx::PgPool;
+    use ulid::Ulid;
 
     use super::UpstreamSessionsCookie;
     use crate::test_utils::{CookieHelper, RequestBuilderExt, ResponseExt, TestState, setup};
@@ -1250,33 +1248,41 @@ mod tests {
         let response = state.request(request).await;
         cookies.save_cookies(&response);
         response.assert_status(StatusCode::SEE_OTHER);
+        let location = response.headers().get(hyper::header::LOCATION).unwrap();
+        // Grab the registration ID from the redirected URL:
+        //   /register/steps/{id}/finish
+        let registration_id: Ulid = str::from_utf8(location.as_bytes())
+            .unwrap()
+            .rsplit('/')
+            .nth(1)
+            .expect("Location to have two slashes")
+            .parse()
+            .expect("last segment of location to be a ULID");
 
         // Check that we have a registered user, with the email imported
         let mut repo = state.repository().await.unwrap();
-        let user = repo
-            .user()
-            .find_by_username("john")
+        let registration: UserRegistration = repo
+            .user_registration()
+            .lookup(registration_id)
             .await
             .unwrap()
-            .expect("user exists");
+            .expect("user registration exists");
 
-        let link = repo
-            .upstream_oauth_link()
-            .find_by_subject(&provider, "subject")
-            .await
-            .unwrap()
-            .expect("link exists");
-
-        assert_eq!(link.user_id, Some(user.id));
+        assert_eq!(registration.password, None);
+        assert_eq!(registration.completed_at, None);
+        assert_eq!(registration.username, "john");
 
-        let page = repo
+        let email_auth_id = registration
+            .email_authentication_id
+            .expect("registration should have an email authentication");
+        let email_auth: UserEmailAuthentication = repo
             .user_email()
-            .list(UserEmailFilter::new().for_user(&user), Pagination::first(1))
+            .lookup_authentication(email_auth_id)
             .await
-            .unwrap();
-        let edge = page.edges.first().expect("email exists");
-
-        assert_eq!(edge.node.email, "[email protected]");
+            .unwrap()
+            .expect("email authentication should exist");
+        assert_eq!(email_auth.email, "[email protected]");
+        assert!(email_auth.completed_at.is_some());
     }
 
     #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]