Graphite/editor/src/test_utils.rs

use crate::application::Editor;
use crate::messages::input_mapper::utility_types::keyboard::ModifierKeys;
use crate::messages::input_mapper::utility_types::pointer::{EditorPointerState, MouseKeys, ViewportPosition};
use crate::messages::portfolio::document::node_graph::document_node_definitions::DefinitionIdentifier;
use crate::messages::prelude::*;
use crate::messages::tool::tool_messages::tool_prelude::Key;
use crate::messages::tool::utility_types::ToolType;
use crate::node_graph_executor::Instrumented;
use crate::node_graph_executor::NodeRuntime;
use crate::test_utils::test_prelude::LayerNodeIdentifier;
use glam::{DVec2, UVec2};
use graphene_std::raster::color::Color;
use graphene_std::uuid::NodeId;

/// A set of utility functions to make the writing of editor test more declarative
pub struct EditorTestUtils {
	pub editor: Editor,
	pub runtime: NodeRuntime,
}

impl EditorTestUtils {
	pub fn create() -> Self {
		let _ = env_logger::builder().is_test(true).try_init();

		let (mut editor, runtime) = Editor::new_local_executor();

		editor.handle_message(PortfolioMessage::Init);

		// Enable the dual-write soak validation in tests so storage round-trip drift fails loud, matching
		// the previous always-on-in-debug behavior now that it is gated by a preference.
		editor.handle_message(PreferencesMessage::ValidateStorageRoundTrip { enabled: true });

		Self { editor, runtime }
	}

	pub fn eval_graph<'a>(&'a mut self) -> impl std::future::Future<Output = Result<Instrumented, String>> + 'a {
		// An inner function is required since async functions in traits are a bit weird
		async fn run<'a>(editor: &'a mut Editor, runtime: &'a mut NodeRuntime) -> Result<Instrumented, String> {
			let portfolio = &mut editor.dispatcher.message_handlers.portfolio_message_handler;
			let document_id = portfolio.active_document_id.unwrap();
			let (executor, documents) = (&mut portfolio.executor, &mut portfolio.documents);
			let document = documents.get_mut(&document_id).unwrap();

			let instrumented = match executor.update_node_graph_instrumented(document) {
				Ok(instrumented) => instrumented,
				Err(e) => return Err(format!("update_node_graph_instrumented failed\n\n{e}")),
			};

			if let Err(e) = executor.submit_current_node_graph_evaluation(document, document_id, UVec2::ONE, 1., Default::default(), DVec2::ZERO) {
				return Err(format!("submit_current_node_graph_evaluation failed\n\n{e}"));
			}
			runtime.run().await;

			// An async source reports `Pending` on the evaluation that starts it and marks the runtime dirty once
			// it completes, so the value only reaches the render on a follow-up evaluation. The superseded response's
			// `Pending` error is ignored, but its messages carry the incremental resolved-types delta and must be
			// dispatched, or the editor's type map desyncs permanently.
			while runtime.take_dirty() {
				let mut superseded_messages = VecDeque::new();
				let _ = editor.poll_node_graph_evaluation(&mut superseded_messages);
				for message in superseded_messages {
					editor.handle_message(message);
				}

				let portfolio = &mut editor.dispatcher.message_handlers.portfolio_message_handler;
				let (executor, documents) = (&mut portfolio.executor, &mut portfolio.documents);
				let document = documents.get_mut(&document_id).unwrap();

				if let Err(e) = executor.submit_current_node_graph_evaluation(document, document_id, UVec2::ONE, 1., Default::default(), DVec2::ZERO) {
					return Err(format!("submit_current_node_graph_evaluation failed\n\n{e}"));
				}
				runtime.run().await;
			}

			let mut messages = VecDeque::new();
			if let Err(e) = editor.poll_node_graph_evaluation(&mut messages) {
				return Err(format!("Graph should render\n\n{e}"));
			}
			let frontend_messages = messages.into_iter().flat_map(|message| editor.handle_message(message));

			for message in frontend_messages {
				message.check_node_graph_error();
			}

			Ok(instrumented)
		}

		run(&mut self.editor, &mut self.runtime)
	}

	pub async fn handle_message(&mut self, message: impl Into<Message>) -> Vec<FrontendMessage> {
		let frontend_messages_from_msg = self.editor.handle_message(message);

		// Required to process any buffered messages
		if let Err(e) = self.eval_graph().await {
			panic!("Failed to evaluate graph: {e}");
		}

		// Sweep the network interface's structural invariants so any desync fails at the message that caused it
		let violations = self.active_document().network_interface.validate_invariants();
		assert!(violations.is_empty(), "Network interface invariant violations:\n{}", violations.join("\n"));

		frontend_messages_from_msg
	}

	pub async fn new_document(&mut self) {
		self.handle_message(Message::Portfolio(PortfolioMessage::NewDocumentWithName { name: String::from("Test document") }))
			.await;
	}

	pub async fn draw_rect(&mut self, x1: f64, y1: f64, x2: f64, y2: f64) {
		self.drag_tool(ToolType::Rectangle, x1, y1, x2, y2, ModifierKeys::default()).await;
	}

	pub async fn draw_polygon(&mut self, x1: f64, y1: f64, x2: f64, y2: f64) {
		self.drag_tool(ToolType::Shape, x1, y1, x2, y2, ModifierKeys::default()).await;
	}

	pub async fn draw_ellipse(&mut self, x1: f64, y1: f64, x2: f64, y2: f64) {
		self.drag_tool(ToolType::Ellipse, x1, y1, x2, y2, ModifierKeys::default()).await;
	}

	pub async fn click_tool(&mut self, typ: ToolType, button: MouseKeys, position: DVec2, modifier_keys: ModifierKeys) {
		self.select_tool(typ).await;

		self.move_mouse(position.x, position.y, modifier_keys, MouseKeys::empty()).await;

		self.mousedown(
			EditorPointerState {
				editor_position: position,
				mouse_keys: button,
				..Default::default()
			},
			modifier_keys,
		)
		.await;

		self.mouseup(
			EditorPointerState {
				editor_position: position,
				..Default::default()
			},
			modifier_keys,
		)
		.await;
	}

	pub async fn drag_tool(&mut self, typ: ToolType, x1: f64, y1: f64, x2: f64, y2: f64, modifier_keys: ModifierKeys) {
		self.select_tool(typ).await;

		self.move_mouse(x1, y1, modifier_keys, MouseKeys::empty()).await;

		self.left_mousedown(x1, y1, modifier_keys).await;

		self.move_mouse(x2, y2, modifier_keys, MouseKeys::LEFT).await;

		self.mouseup(
			EditorPointerState {
				editor_position: (x2, y2).into(),
				mouse_keys: MouseKeys::empty(),
				..Default::default()
			},
			modifier_keys,
		)
		.await;
	}

	pub async fn drag_tool_cancel_rmb(&mut self, typ: ToolType) {
		self.select_tool(typ).await;

		self.move_mouse(50., 50., ModifierKeys::default(), MouseKeys::empty()).await;

		self.left_mousedown(50., 50., ModifierKeys::default()).await;

		self.move_mouse(100., 100., ModifierKeys::default(), MouseKeys::LEFT).await;

		self.mousedown(
			EditorPointerState {
				editor_position: (100., 100.).into(),
				mouse_keys: MouseKeys::LEFT | MouseKeys::RIGHT,
				..Default::default()
			},
			ModifierKeys::default(),
		)
		.await;
	}

	pub fn active_document(&self) -> &DocumentMessageHandler {
		self.editor.dispatcher.message_handlers.portfolio_message_handler.active_document().unwrap()
	}

	pub fn active_document_mut(&mut self) -> &mut DocumentMessageHandler {
		self.editor.dispatcher.message_handlers.portfolio_message_handler.active_document_mut().unwrap()
	}

	pub async fn move_mouse(&mut self, x: f64, y: f64, modifier_keys: ModifierKeys, mouse_keys: MouseKeys) {
		let editor_mouse_state = EditorPointerState {
			editor_position: ViewportPosition::new(x, y),
			mouse_keys,
			..Default::default()
		};
		self.input(InputPreprocessorMessage::PointerMove { editor_mouse_state, modifier_keys }).await;
	}

	pub async fn mousedown(&mut self, editor_mouse_state: EditorPointerState, modifier_keys: ModifierKeys) {
		self.input(InputPreprocessorMessage::PointerDown { editor_mouse_state, modifier_keys }).await;
	}

	pub async fn mouseup(&mut self, editor_mouse_state: EditorPointerState, modifier_keys: ModifierKeys) {
		self.handle_message(InputPreprocessorMessage::PointerUp { editor_mouse_state, modifier_keys }).await;
	}

	pub async fn press(&mut self, key: Key, modifier_keys: ModifierKeys) {
		let key_repeat = false;

		self.handle_message(InputPreprocessorMessage::KeyDown { key, modifier_keys, key_repeat }).await;
		self.handle_message(InputPreprocessorMessage::KeyUp { key, modifier_keys, key_repeat }).await;
	}

	pub async fn left_mousedown(&mut self, x: f64, y: f64, modifier_keys: ModifierKeys) {
		self.mousedown(
			EditorPointerState {
				editor_position: (x, y).into(),
				mouse_keys: MouseKeys::LEFT,
				..Default::default()
			},
			modifier_keys,
		)
		.await;
	}

	pub async fn left_mouseup(&mut self, x: f64, y: f64, modifier_keys: ModifierKeys) {
		self.mouseup(
			EditorPointerState {
				editor_position: (x, y).into(),
				mouse_keys: MouseKeys::empty(),
				..Default::default()
			},
			modifier_keys,
		)
		.await;
	}

	pub async fn input(&mut self, message: InputPreprocessorMessage) {
		self.handle_message(Message::InputPreprocessor(message)).await;
	}

	pub async fn select_tool(&mut self, tool_type: ToolType) {
		match tool_type {
			ToolType::Line => self.handle_message(Message::Tool(ToolMessage::ActivateToolShapeLine)).await,
			ToolType::Rectangle => self.handle_message(Message::Tool(ToolMessage::ActivateToolShapeRectangle)).await,
			ToolType::Ellipse => self.handle_message(Message::Tool(ToolMessage::ActivateToolShapeEllipse)).await,
			_ => self.handle_message(Message::Tool(ToolMessage::ActivateTool { tool_type })).await,
		};
	}

	pub async fn select_primary_color(&mut self, color: Color) {
		self.handle_message(Message::Tool(ToolMessage::SelectWorkingColor { color, primary: true })).await;
	}

	pub async fn select_secondary_color(&mut self, color: Color) {
		self.handle_message(Message::Tool(ToolMessage::SelectWorkingColor { color, primary: false })).await;
	}

	pub async fn create_raster_image(&mut self, image: graphene_std::raster::Image<Color>, mouse: Option<(f64, f64)>) {
		self.handle_message(PortfolioMessage::InsertImage {
			name: None,
			image,
			mouse,
			parent_and_insert_index: None,
		})
		.await;
	}
	pub async fn draw_spline(&mut self, points: &[DVec2]) {
		self.select_tool(ToolType::Spline).await;

		for &point in points {
			self.click_tool(ToolType::Spline, MouseKeys::LEFT, point, ModifierKeys::empty()).await;
		}

		self.press(Key::Enter, ModifierKeys::empty()).await;
	}

	pub async fn get_selected_layer(&mut self) -> Option<LayerNodeIdentifier> {
		self.active_document().network_interface.selected_nodes().selected_layers(self.active_document().metadata()).next()
	}

	pub async fn double_click(&mut self, position: DVec2) {
		self.handle_message(InputPreprocessorMessage::DoubleClick {
			editor_mouse_state: EditorPointerState {
				editor_position: position,
				mouse_keys: MouseKeys::LEFT,
				..Default::default()
			},
			modifier_keys: ModifierKeys::empty(),
		})
		.await;
	}

	pub async fn drag_path(&mut self, points: &[DVec2], modifier_keys: ModifierKeys) {
		if points.is_empty() {
			return;
		}

		let first_point = points[0];
		self.move_mouse(first_point.x, first_point.y, modifier_keys, MouseKeys::empty()).await;
		self.left_mousedown(first_point.x, first_point.y, modifier_keys).await;

		for &point in &points[1..] {
			self.move_mouse(point.x, point.y, modifier_keys, MouseKeys::LEFT).await;
		}

		self.mouseup(
			EditorPointerState {
				editor_position: points[points.len() - 1],
				mouse_keys: MouseKeys::empty(),
				..Default::default()
			},
			modifier_keys,
		)
		.await;
	}

	/// Necessary for doing snapping since snaps outside of the viewport are discarded
	pub async fn set_viewport_size(&mut self, top_left: DVec2, bottom_right: DVec2) {
		self.handle_message(ViewportMessage::Update {
			x: top_left.x,
			y: top_left.y,
			width: bottom_right.x - top_left.x,
			height: bottom_right.y - top_left.y,
			scale: 1.,
		})
		.await;
	}

	pub async fn create_node_by_name(&mut self, node_type: DefinitionIdentifier) -> NodeId {
		self.create_node(node_type, None).await
	}

	pub async fn create_node_by_name_at(&mut self, node_type: DefinitionIdentifier, x: i32, y: i32) -> NodeId {
		self.create_node(node_type, Some((x, y))).await
	}

	async fn create_node(&mut self, node_type: DefinitionIdentifier, xy: Option<(i32, i32)>) -> NodeId {
		let node_id = NodeId::new();
		self.handle_message(NodeGraphMessage::CreateNodeFromContextMenu {
			node_id: Some(node_id),
			node_type,
			xy,
			add_transaction: true,
		})
		.await;
		node_id
	}
}

pub trait FrontendMessageTestUtils {
	fn check_node_graph_error(&self);
}

impl FrontendMessageTestUtils for FrontendMessage {
	fn check_node_graph_error(&self) {
		let FrontendMessage::UpdateNodeGraphErrorDiagnostic { error } = self else { return };
		if let Some(error) = error {
			panic!("error: {:?}", error);
		}
	}
}

#[cfg(test)]
pub mod test_prelude {
	pub use super::FrontendMessageTestUtils;
	pub use crate::application::Editor;
	pub use crate::float_eq;
	pub use crate::messages::clipboard::utility_types::ClipboardContentRaw;
	pub use crate::messages::input_mapper::utility_types::keyboard::{Key, ModifierKeys};
	pub use crate::messages::input_mapper::utility_types::pointer::MouseKeys;
	pub use crate::messages::portfolio::document::node_graph::document_node_definitions::DefinitionIdentifier;
	pub use crate::messages::portfolio::document::utility_types::document_metadata::LayerNodeIdentifier;
	pub use crate::messages::prelude::*;
	pub use crate::messages::tool::common_functionality::graph_modification_utils::{NodeGraphLayer, is_layer_fed_by_node_of_name};
	pub use crate::messages::tool::utility_types::ToolType;
	pub use crate::node_graph_executor::NodeRuntime;
	pub use crate::test_utils::EditorTestUtils;
	pub use core::f64;
	pub use glam::{DVec2, IVec2};
	pub use graph_craft::document::DocumentNode;
	pub use graphene_std::raster::{Color, Image};
	pub use graphene_std::transform::Footprint;

	#[macro_export]
	macro_rules! float_eq {
		($left:expr, $right:expr $(,)?) => {
			match (&$left, &$right) {
				(left_val, right_val) => {
					if (*left_val - *right_val).abs() > 1e-10 {
						panic!("assertion `left == right` failed\n  left: {}\n right: {}", *left_val, *right_val)
					}
				}
			}
		};
	}
}