feat: Tokenserver: Add node assignment logic

src/db/error.rs

@@ -38,8 +38,8 @@ pub enum DbErrorKind {
     #[error("Specified batch does not exist")]
     BatchNotFound,
 
-    #[error("Tokenserver user not found")]
-    TokenserverUserNotFound,
+    #[error("Tokenserver user retired")]
+    TokenserverUserRetired,
 
     #[error("An attempt at a conflicting write")]
     Conflict,
@@ -84,9 +84,7 @@ impl DbError {
 impl From<DbErrorKind> for DbError {
     fn from(kind: DbErrorKind) -> Self {
         let status = match kind {
-            DbErrorKind::TokenserverUserNotFound
-            | DbErrorKind::CollectionNotFound
-            | DbErrorKind::BsoNotFound => StatusCode::NOT_FOUND,
+            DbErrorKind::CollectionNotFound | DbErrorKind::BsoNotFound => StatusCode::NOT_FOUND,
             // Matching the Python code here (a 400 vs 404)
             DbErrorKind::BatchNotFound | DbErrorKind::SpannerTooLarge(_) => StatusCode::BAD_REQUEST,
             // NOTE: the protocol specification states that we should return a
@@ -96,6 +94,8 @@ impl From<DbErrorKind> for DbError {
             //  * android bug: https://bugzilla.mozilla.org/show_bug.cgi?id=959032
             DbErrorKind::Conflict => StatusCode::SERVICE_UNAVAILABLE,
             DbErrorKind::Quota => StatusCode::FORBIDDEN,
+            // NOTE: TokenserverUserRetired is an internal service error for compatibility reasons
+            // (the legacy Tokenserver returned an internal service error in this situation)
             _ => StatusCode::INTERNAL_SERVER_ERROR,
         };
 

src/tokenserver/db/mock.rs

@@ -37,8 +37,8 @@ impl MockDb {
 }
 
 impl Db for MockDb {
-    fn get_user(&self, _params: params::GetUser) -> DbFuture<'_, results::GetUser> {
-        Box::pin(future::ok(results::GetUser::default()))
+    fn replace_user(&self, _params: params::ReplaceUser) -> DbFuture<'_, results::ReplaceUser> {
+        Box::pin(future::ok(()))
     }
 
     fn replace_users(&self, _params: params::ReplaceUsers) -> DbFuture<'_, results::ReplaceUsers> {
@@ -49,6 +49,10 @@ impl Db for MockDb {
         Box::pin(future::ok(results::PostUser::default()))
     }
 
+    fn allocate_user(&self, _params: params::AllocateUser) -> DbFuture<'_, results::AllocateUser> {
+        Box::pin(future::ok(results::AllocateUser::default()))
+    }
+
     fn put_user(&self, _params: params::PutUser) -> DbFuture<'_, results::PutUser> {
         Box::pin(future::ok(()))
     }
@@ -61,6 +65,24 @@ impl Db for MockDb {
         Box::pin(future::ok(results::GetNodeId::default()))
     }
 
+    fn get_best_node(&self, _params: params::GetBestNode) -> DbFuture<'_, results::GetBestNode> {
+        Box::pin(future::ok(results::GetBestNode::default()))
+    }
+
+    fn add_user_to_node(
+        &self,
+        _params: params::AddUserToNode,
+    ) -> DbFuture<'_, results::AddUserToNode> {
+        Box::pin(future::ok(()))
+    }
+
+    fn get_or_create_user(
+        &self,
+        _params: params::GetOrCreateUser,
+    ) -> DbFuture<'_, results::GetOrCreateUser> {
+        Box::pin(future::ok(results::GetOrCreateUser::default()))
+    }
+
     #[cfg(test)]
     fn set_user_created_at(
         &self,
@@ -69,6 +91,19 @@ impl Db for MockDb {
         Box::pin(future::ok(()))
     }
 
+    #[cfg(test)]
+    fn set_user_replaced_at(
+        &self,
+        _params: params::SetUserReplacedAt,
+    ) -> DbFuture<'_, results::SetUserReplacedAt> {
+        Box::pin(future::ok(()))
+    }
+
+    #[cfg(test)]
+    fn get_user(&self, _params: params::GetUser) -> DbFuture<'_, results::GetUser> {
+        Box::pin(future::ok(results::GetUser::default()))
+    }
+
     #[cfg(test)]
     fn get_users(&self, _params: params::GetRawUsers) -> DbFuture<'_, results::GetRawUsers> {
         Box::pin(future::ok(results::GetRawUsers::default()))
@@ -79,6 +114,21 @@ impl Db for MockDb {
         Box::pin(future::ok(results::PostNode::default()))
     }
 
+    #[cfg(test)]
+    fn get_node(&self, _params: params::GetNode) -> DbFuture<'_, results::GetNode> {
+        Box::pin(future::ok(results::GetNode::default()))
+    }
+
+    #[cfg(test)]
+    fn unassign_node(&self, _params: params::UnassignNode) -> DbFuture<'_, results::UnassignNode> {
+        Box::pin(future::ok(()))
+    }
+
+    #[cfg(test)]
+    fn remove_node(&self, _params: params::RemoveNode) -> DbFuture<'_, results::RemoveNode> {
+        Box::pin(future::ok(()))
+    }
+
     #[cfg(test)]
     fn post_service(&self, _params: params::PostService) -> DbFuture<'_, results::PostService> {
         Box::pin(future::ok(results::PostService::default()))