ComfyUI/tests-ui/tests/users.test.js

// @ts-check
/// <reference path="../node_modules/@types/jest/index.d.ts" />
const { start } = require("../utils");
const lg = require("../utils/litegraph");

describe("users", () => {
	beforeEach(() => {
		lg.setup(global);
	});

	afterEach(() => {
		lg.teardown(global);
	});

	function expectNoUserScreen() {
		// Ensure login isnt visible
		const selection = document.querySelectorAll("#comfy-user-selection")?.[0];
		expect(selection["style"].display).toBe("none");
		const menu = document.querySelectorAll(".comfy-menu")?.[0];
		expect(window.getComputedStyle(menu)?.display).not.toBe("none");
	}

	describe("multi-user", () => {
		async function testUserScreen(onShown, users) {
			if (!users) {
				users = {};
			}
			const starting = start({
				resetEnv: true,
				users,
			});

			// Ensure no current user
			expect(localStorage["Comfy.userId"]).toBeFalsy();
			expect(localStorage["Comfy.userName"]).toBeFalsy();

			await new Promise(process.nextTick); // wait for promises to resolve

			const selection = document.querySelectorAll("#comfy-user-selection")?.[0];
			expect(selection).toBeTruthy();

			// Ensure login is visible
			expect(window.getComputedStyle(selection)?.display).not.toBe("none");
			// Ensure menu is hidden
			const menu = document.querySelectorAll(".comfy-menu")?.[0];
			expect(window.getComputedStyle(menu)?.display).toBe("none");

			const isCreate = await onShown(selection);

			// Submit form
			selection.querySelectorAll("form")[0].submit();
			await new Promise(process.nextTick); // wait for promises to resolve

			// Wait for start
			const s = await starting;

			// Ensure login is removed
			expect(document.querySelectorAll("#comfy-user-selection")).toHaveLength(0);
			expect(window.getComputedStyle(menu)?.display).not.toBe("none");

			// Ensure settings + templates are saved
			const { api } = require("../../web/scripts/api");
			expect(api.createUser).toHaveBeenCalledTimes(+isCreate);
			expect(api.storeSettings).toHaveBeenCalledTimes(+isCreate);
			expect(api.storeUserData).toHaveBeenCalledTimes(+isCreate);
			if (isCreate) {
				expect(api.storeUserData).toHaveBeenCalledWith("comfy.templates.json", null, { stringify: false });
				expect(s.app.isNewUserSession).toBeTruthy();
			} else {
				expect(s.app.isNewUserSession).toBeFalsy();
			}

			return { users, selection, ...s };
		}

		it("allows user creation if no users", async () => {
			const { users } = await testUserScreen((selection) => {
				// Ensure we have no users flag added
				expect(selection.classList.contains("no-users")).toBeTruthy();

				// Enter a username
				const input = selection.getElementsByTagName("input")[0];
				input.focus();
				input.value = "Test User";

				return true;
			});

			expect(users).toStrictEqual({
				"Test User!": "Test User",
			});

			expect(localStorage["Comfy.userId"]).toBe("Test User!");
			expect(localStorage["Comfy.userName"]).toBe("Test User");
		});
		it("allows user creation if no current user but other users", async () => {
			const users = {
				"Test User 2!": "Test User 2",
			};

			await testUserScreen((selection) => {
				expect(selection.classList.contains("no-users")).toBeFalsy();

				// Enter a username
				const input = selection.getElementsByTagName("input")[0];
				input.focus();
				input.value = "Test User 3";
				return true;
			}, users);

			expect(users).toStrictEqual({
				"Test User 2!": "Test User 2",
				"Test User 3!": "Test User 3",
			});

			expect(localStorage["Comfy.userId"]).toBe("Test User 3!");
			expect(localStorage["Comfy.userName"]).toBe("Test User 3");
		});
		it("allows user selection if no current user but other users", async () => {
			const users = {
				"A!": "A",
				"B!": "B",
				"C!": "C",
			};

			await testUserScreen((selection) => {
				expect(selection.classList.contains("no-users")).toBeFalsy();

				// Check user list
				const select = selection.getElementsByTagName("select")[0];
				const options = select.getElementsByTagName("option");
				expect(
					[...options]
						.filter((o) => !o.disabled)
						.reduce((p, n) => {
							p[n.getAttribute("value")] = n.textContent;
							return p;
						}, {})
				).toStrictEqual(users);

				// Select an option
				select.focus();
				select.value = options[2].value;

				return false;
			}, users);

			expect(users).toStrictEqual(users);

			expect(localStorage["Comfy.userId"]).toBe("B!");
			expect(localStorage["Comfy.userName"]).toBe("B");
		});
		it("doesnt show user screen if current user", async () => {
			const starting = start({
				resetEnv: true,
				users: {
					"User!": "User",
				},
				localStorage: {
					"Comfy.userId": "User!",
					"Comfy.userName": "User",
				},
			});
			await new Promise(process.nextTick); // wait for promises to resolve

			expectNoUserScreen();

			await starting;
		});
		it("allows user switching", async () => {
			const { app } = await start({
				resetEnv: true,
				users: {
					"User!": "User",
				},
				localStorage: {
					"Comfy.userId": "User!",
					"Comfy.userName": "User",
				},
			});

			// cant actually test switching user easily but can check the setting is present
			expect(app.ui.settings.settingsLookup["Comfy.SwitchUser"]).toBeTruthy();
		});
	});
	describe("single-user", () => {
		it("doesnt show user creation if no default user", async () => {
			const { app } = await start({
				resetEnv: true,
				users: false,
			});
            expectNoUserScreen();

			// It should store the settings
			const { api } = require("../../web/scripts/api");
			expect(api.storeSettings).toHaveBeenCalledTimes(1);
			expect(api.storeUserData).toHaveBeenCalledTimes(1);
			expect(api.storeUserData).toHaveBeenCalledWith("comfy.templates.json", null, { stringify: false });
			expect(app.isNewUserSession).toBeTruthy();
		});
		it("doesnt show user creation if default user", async () => {
			const { app } = await start({
				resetEnv: true,
				users: true,
			});
            expectNoUserScreen();

			// It should store the settings
			const { api } = require("../../web/scripts/api");
			expect(api.storeSettings).toHaveBeenCalledTimes(0);
			expect(api.storeUserData).toHaveBeenCalledTimes(0);
			expect(app.isNewUserSession).toBeFalsy();
		});
		it("doesnt allow user switching", async () => {
			const { app } = await start({
				resetEnv: true,
			});
            expectNoUserScreen();

			expect(app.ui.settings.settingsLookup["Comfy.SwitchUser"]).toBeFalsy();
		});
	});
});