continuwuity/src/service/rooms/roles/cache_tests.rs

use std::collections::{BTreeMap, HashMap, HashSet};

use conduwuit_core::matrix::space_roles::RoleDefinition;
use ruma::{OwnedRoomId, OwnedUserId, room_id, user_id};

use super::{
	compute_user_power_level, roles_satisfy_requirements,
	tests::{make_roles, make_set},
};

struct MockCache {
	roles: HashMap<OwnedRoomId, BTreeMap<String, RoleDefinition>>,
	user_roles: HashMap<OwnedRoomId, HashMap<OwnedUserId, HashSet<String>>>,
	room_requirements: HashMap<OwnedRoomId, HashMap<OwnedRoomId, HashSet<String>>>,
	room_to_space: HashMap<OwnedRoomId, HashSet<OwnedRoomId>>,
}

impl MockCache {
	fn new() -> Self {
		Self {
			roles: HashMap::new(),
			user_roles: HashMap::new(),
			room_requirements: HashMap::new(),
			room_to_space: HashMap::new(),
		}
	}

	fn add_space(&mut self, space: OwnedRoomId, roles: BTreeMap<String, RoleDefinition>) {
		self.roles.insert(space, roles);
	}

	fn add_child(&mut self, space: &OwnedRoomId, child: OwnedRoomId) {
		self.room_to_space
			.entry(child)
			.or_default()
			.insert(space.clone());
	}

	fn assign_role(&mut self, space: &OwnedRoomId, user: OwnedUserId, role: String) {
		self.user_roles
			.entry(space.clone())
			.or_default()
			.entry(user)
			.or_default()
			.insert(role);
	}

	fn revoke_role(&mut self, space: &OwnedRoomId, user: &OwnedUserId, role: &str) {
		if let Some(space_users) = self.user_roles.get_mut(space) {
			if let Some(user_roles) = space_users.get_mut(user) {
				user_roles.remove(role);
			}
		}
	}

	fn set_room_requirements(
		&mut self,
		space: &OwnedRoomId,
		room: OwnedRoomId,
		reqs: HashSet<String>,
	) {
		self.room_requirements
			.entry(space.clone())
			.or_default()
			.insert(room, reqs);
	}

	fn user_qualifies(
		&self,
		space: &OwnedRoomId,
		room: &OwnedRoomId,
		user: &OwnedUserId,
	) -> bool {
		let Some(required) = self.room_requirements.get(space).and_then(|r| r.get(room)) else {
			return true;
		};
		if required.is_empty() {
			return true;
		}
		let Some(assigned) = self.user_roles.get(space).and_then(|u| u.get(user)) else {
			return false;
		};
		roles_satisfy_requirements(required, assigned)
	}

	fn get_power_level(&self, space: &OwnedRoomId, user: &OwnedUserId) -> Option<i64> {
		let role_defs = self.roles.get(space)?;
		let assigned = self.user_roles.get(space)?.get(user)?;
		compute_user_power_level(role_defs, assigned)
	}
}

#[test]
fn cache_populate_and_lookup() {
	let mut cache = MockCache::new();
	let space = room_id!("!space:example.com").to_owned();
	let child = room_id!("!child:example.com").to_owned();
	let alice = user_id!("@alice:example.com").to_owned();

	cache.add_space(space.clone(), make_roles(&[("admin", Some(100)), ("nsfw", None)]));
	cache.add_child(&space, child.clone());
	cache.assign_role(&space, alice.clone(), "nsfw".to_owned());
	cache.set_room_requirements(&space, child.clone(), make_set(&["nsfw"]));

	assert!(cache.user_qualifies(&space, &child, &alice));
	assert_eq!(cache.get_power_level(&space, &alice), None);
}

#[test]
fn cache_invalidation_on_role_revoke() {
	let mut cache = MockCache::new();
	let space = room_id!("!space:example.com").to_owned();
	let child = room_id!("!nsfw:example.com").to_owned();
	let alice = user_id!("@alice:example.com").to_owned();

	cache.add_space(space.clone(), make_roles(&[("nsfw", None)]));
	cache.assign_role(&space, alice.clone(), "nsfw".to_owned());
	cache.set_room_requirements(&space, child.clone(), make_set(&["nsfw"]));

	assert!(cache.user_qualifies(&space, &child, &alice));

	cache.revoke_role(&space, &alice, "nsfw");
	assert!(!cache.user_qualifies(&space, &child, &alice));
}

#[test]
fn cache_invalidation_on_requirement_change() {
	let mut cache = MockCache::new();
	let space = room_id!("!space:example.com").to_owned();
	let child = room_id!("!room:example.com").to_owned();
	let alice = user_id!("@alice:example.com").to_owned();

	cache.add_space(space.clone(), make_roles(&[("nsfw", None), ("vip", None)]));
	cache.assign_role(&space, alice.clone(), "vip".to_owned());
	cache.set_room_requirements(&space, child.clone(), make_set(&["vip"]));

	assert!(cache.user_qualifies(&space, &child, &alice));

	cache.set_room_requirements(&space, child.clone(), make_set(&["vip", "nsfw"]));
	assert!(!cache.user_qualifies(&space, &child, &alice));
}

#[test]
fn cache_power_level_updates_on_role_change() {
	let mut cache = MockCache::new();
	let space = room_id!("!space:example.com").to_owned();
	let alice = user_id!("@alice:example.com").to_owned();

	cache.add_space(space.clone(), make_roles(&[("admin", Some(100)), ("mod", Some(50))]));

	assert_eq!(cache.get_power_level(&space, &alice), None);

	cache.assign_role(&space, alice.clone(), "mod".to_owned());
	assert_eq!(cache.get_power_level(&space, &alice), Some(50));

	cache.assign_role(&space, alice.clone(), "admin".to_owned());
	assert_eq!(cache.get_power_level(&space, &alice), Some(100));

	cache.revoke_role(&space, &alice, "admin");
	assert_eq!(cache.get_power_level(&space, &alice), Some(50));
}