feat: add macros for testing (#2337)

* feat: add get_models_test_class_hashes macro

* feat: add new spawn_test_world_full macro

* fix: enhance macro for spawn world

crates/dojo-lang/src/inline_macros/get_models_test_class_hashes.rs

@@ -0,0 +1,104 @@
+use cairo_lang_defs::patcher::PatchBuilder;
+use cairo_lang_defs::plugin::{
+    InlineMacroExprPlugin, InlinePluginResult, MacroPluginMetadata, NamedPlugin, PluginDiagnostic,
+    PluginGeneratedFile,
+};
+use cairo_lang_defs::plugin_utils::unsupported_bracket_diagnostic;
+use cairo_lang_diagnostics::Severity;
+use cairo_lang_syntax::node::{ast, TypedStablePtr, TypedSyntaxNode};
+
+use super::unsupported_arg_diagnostic;
+use super::utils::{extract_namespaces, load_manifest_models_and_namespaces};
+
+#[derive(Debug, Default)]
+pub struct GetModelsTestClassHashes;
+
+impl NamedPlugin for GetModelsTestClassHashes {
+    const NAME: &'static str = "get_models_test_class_hashes";
+}
+
+impl InlineMacroExprPlugin for GetModelsTestClassHashes {
+    fn generate_code(
+        &self,
+        db: &dyn cairo_lang_syntax::node::db::SyntaxGroup,
+        syntax: &ast::ExprInlineMacro,
+        metadata: &MacroPluginMetadata<'_>,
+    ) -> InlinePluginResult {
+        let ast::WrappedArgList::ParenthesizedArgList(arg_list) = syntax.arguments(db) else {
+            return unsupported_bracket_diagnostic(db, syntax);
+        };
+
+        let args = arg_list.arguments(db).elements(db);
+
+        if args.len() > 1 {
+            return InlinePluginResult {
+                code: None,
+                diagnostics: vec![PluginDiagnostic {
+                    stable_ptr: syntax.stable_ptr().untyped(),
+                    message: "Invalid arguments. Expected \
+                              \"get_models_test_class_hashes!([\"ns1\", \"ns2\")]\" or \
+                              \"get_models_test_class_hashes!()\"."
+                        .to_string(),
+                    severity: Severity::Error,
+                }],
+            };
+        }
+
+        let whitelisted_namespaces = if args.len() == 1 {
+            let ast::ArgClause::Unnamed(expected_array) = args[0].arg_clause(db) else {
+                return unsupported_arg_diagnostic(db, syntax);
+            };
+
+            match extract_namespaces(db, &expected_array.value(db)) {
+                Ok(namespaces) => namespaces,
+                Err(e) => {
+                    return InlinePluginResult { code: None, diagnostics: vec![e] };
+                }
+            }
+        } else {
+            vec![]
+        };
+
+        let (_namespaces, models) =
+            match load_manifest_models_and_namespaces(metadata.cfg_set, &whitelisted_namespaces) {
+                Ok((namespaces, models)) => (namespaces, models),
+                Err(_e) => {
+                    return InlinePluginResult {
+                        code: None,
+                        diagnostics: vec![PluginDiagnostic {
+                            stable_ptr: syntax.stable_ptr().untyped(),
+                            message: "Failed to load models and namespaces, ensure you have run \
+                                      `sozo build` first."
+                                .to_string(),
+                            severity: Severity::Error,
+                        }],
+                    };
+                }
+            };
+
+        let mut builder = PatchBuilder::new(db, syntax);
+
+        // Use the TEST_CLASS_HASH for each model, which is already a qualified path, no `use`
+        // required.
+        builder.add_str(&format!(
+            "[{}].span()",
+            models
+                .iter()
+                .map(|m| format!("{}::TEST_CLASS_HASH", m))
+                .collect::<Vec<String>>()
+                .join(", ")
+        ));
+
+        let (code, code_mappings) = builder.build();
+
+        InlinePluginResult {
+            code: Some(PluginGeneratedFile {
+                name: "get_models_test_class_hashes_macro".into(),
+                content: code,
+                code_mappings,
+                aux_data: None,
+            }),
+            diagnostics: vec![],
+        }
+    }
+}

crates/dojo-lang/src/inline_macros/mod.rs

@@ -7,8 +7,10 @@ use smol_str::SmolStr;
 pub mod delete;
 pub mod emit;
 pub mod get;
+pub mod get_models_test_class_hashes;
 pub mod selector_from_tag;
 pub mod set;
+pub mod spawn_test_world;
 pub mod utils;
 
 const CAIRO_ERR_MSG_LEN: usize = 31;

crates/dojo-lang/src/inline_macros/spawn_test_world.rs

@@ -0,0 +1,102 @@
+use cairo_lang_defs::patcher::PatchBuilder;
+use cairo_lang_defs::plugin::{
+    InlineMacroExprPlugin, InlinePluginResult, MacroPluginMetadata, NamedPlugin, PluginDiagnostic,
+    PluginGeneratedFile,
+};
+use cairo_lang_defs::plugin_utils::unsupported_bracket_diagnostic;
+use cairo_lang_diagnostics::Severity;
+use cairo_lang_syntax::node::{ast, TypedStablePtr, TypedSyntaxNode};
+
+use super::unsupported_arg_diagnostic;
+use super::utils::{extract_namespaces, load_manifest_models_and_namespaces};
+
+#[derive(Debug, Default)]
+pub struct SpawnTestWorld;
+
+impl NamedPlugin for SpawnTestWorld {
+    const NAME: &'static str = "spawn_test_world";
+}
+
+impl InlineMacroExprPlugin for SpawnTestWorld {
+    fn generate_code(
+        &self,
+        db: &dyn cairo_lang_syntax::node::db::SyntaxGroup,
+        syntax: &ast::ExprInlineMacro,
+        metadata: &MacroPluginMetadata<'_>,
+    ) -> InlinePluginResult {
+        let ast::WrappedArgList::ParenthesizedArgList(arg_list) = syntax.arguments(db) else {
+            return unsupported_bracket_diagnostic(db, syntax);
+        };
+
+        let args = arg_list.arguments(db).elements(db);
+
+        if args.len() > 1 {
+            return InlinePluginResult {
+                code: None,
+                diagnostics: vec![PluginDiagnostic {
+                    stable_ptr: syntax.stable_ptr().untyped(),
+                    message: "Invalid arguments. Expected \"spawn_test_world!()\" or \
+                              \"spawn_test_world!([\"ns1\"])"
+                        .to_string(),
+                    severity: Severity::Error,
+                }],
+            };
+        }
+
+        let whitelisted_namespaces = if args.len() == 1 {
+            let ast::ArgClause::Unnamed(expected_array) = args[0].arg_clause(db) else {
+                return unsupported_arg_diagnostic(db, syntax);
+            };
+
+            match extract_namespaces(db, &expected_array.value(db)) {
+                Ok(namespaces) => namespaces,
+                Err(e) => {
+                    return InlinePluginResult { code: None, diagnostics: vec![e] };
+                }
+            }
+        } else {
+            vec![]
+        };
+
+        let (namespaces, models) =
+            match load_manifest_models_and_namespaces(metadata.cfg_set, &whitelisted_namespaces) {
+                Ok((namespaces, models)) => (namespaces, models),
+                Err(_e) => {
+                    return InlinePluginResult {
+                        code: None,
+                        diagnostics: vec![PluginDiagnostic {
+                            stable_ptr: syntax.stable_ptr().untyped(),
+                            message: "Failed to load models and namespaces, ensure you have run \
+                                      `sozo build` first."
+                                .to_string(),
+                            severity: Severity::Error,
+                        }],
+                    };
+                }
+            };
+
+        let mut builder = PatchBuilder::new(db, syntax);
+
+        builder.add_str(&format!(
+            "dojo::utils::test::spawn_test_world([{}].span(), [{}].span())",
+            namespaces.iter().map(|n| format!("\"{}\"", n)).collect::<Vec<String>>().join(", "),
+            models
+                .iter()
+                .map(|m| format!("{}::TEST_CLASS_HASH", m))
+                .collect::<Vec<String>>()
+                .join(", ")
+        ));
+
+        let (code, code_mappings) = builder.build();
+
+        InlinePluginResult {
+            code: Some(PluginGeneratedFile {
+                name: "spawn_test_world_macro".into(),
+                content: code,
+                code_mappings,
+                aux_data: None,
+            }),
+            diagnostics: vec![],
+        }
+    }
+}

crates/dojo-lang/src/inline_macros/utils.rs

@@ -1,6 +1,16 @@
-use cairo_lang_syntax::node::ast::{ExprPath, ExprStructCtorCall};
+use std::collections::HashSet;
+
+use cairo_lang_defs::plugin::PluginDiagnostic;
+use cairo_lang_diagnostics::Severity;
+use cairo_lang_filesystem::cfg::CfgSet;
+use cairo_lang_syntax::node::ast::{self, ExprPath, ExprStructCtorCall};
+use cairo_lang_syntax::node::db::SyntaxGroup;
 use cairo_lang_syntax::node::kind::SyntaxKind;
-use cairo_lang_syntax::node::SyntaxNode;
+use cairo_lang_syntax::node::{SyntaxNode, TypedStablePtr, TypedSyntaxNode};
+use camino::Utf8PathBuf;
+use dojo_world::config::namespace_config::DOJO_MANIFESTS_DIR_CFG_KEY;
+use dojo_world::contracts::naming;
+use dojo_world::manifest::BaseManifest;
 
 #[derive(Debug)]
 pub enum SystemRWOpRecord {
@@ -22,3 +32,79 @@ pub fn parent_of_kind(
     }
     None
 }
+
+/// Reads all the models and namespaces from base manifests files.
+pub fn load_manifest_models_and_namespaces(
+    cfg_set: &CfgSet,
+    whitelisted_namespaces: &[String],
+) -> anyhow::Result<(Vec<String>, Vec<String>)> {
+    let dojo_manifests_dir = get_dojo_manifests_dir(cfg_set.clone())?;
+
+    let base_dir = dojo_manifests_dir.join("base");
+    let base_abstract_manifest = BaseManifest::load_from_path(&base_dir)?;
+
+    let mut models = HashSet::new();
+    let mut namespaces = HashSet::new();
+
+    for model in base_abstract_manifest.models {
+        let qualified_path = model.inner.qualified_path;
+        let namespace = naming::split_tag(&model.inner.tag)?.0;
+
+        if !whitelisted_namespaces.is_empty() && !whitelisted_namespaces.contains(&namespace) {
+            continue;
+        }
+
+        models.insert(qualified_path);
+        namespaces.insert(namespace);
+    }
+
+    let models_vec: Vec<String> = models.into_iter().collect();
+    let namespaces_vec: Vec<String> = namespaces.into_iter().collect();
+
+    Ok((namespaces_vec, models_vec))
+}
+
+/// Gets the dojo_manifests_dir from the cfg_set.
+pub fn get_dojo_manifests_dir(cfg_set: CfgSet) -> anyhow::Result<Utf8PathBuf> {
+    for cfg in cfg_set.into_iter() {
+        if cfg.key == DOJO_MANIFESTS_DIR_CFG_KEY {
+            return Ok(Utf8PathBuf::from(cfg.value.unwrap().as_str().to_string()));
+        }
+    }
+
+    Err(anyhow::anyhow!("dojo_manifests_dir not found"))
+}
+
+/// Extracts the namespaces from a fixed size array of strings.
+pub fn extract_namespaces(
+    db: &dyn SyntaxGroup,
+    expression: &ast::Expr,
+) -> Result<Vec<String>, PluginDiagnostic> {
+    let mut namespaces = vec![];
+
+    match expression {
+        ast::Expr::FixedSizeArray(array) => {
+            for element in array.exprs(db).elements(db) {
+                if let ast::Expr::String(string_literal) = element {
+                    namespaces.push(string_literal.as_syntax_node().get_text(db).replace('\"', ""));
+                } else {
+                    return Err(PluginDiagnostic {
+                        stable_ptr: element.stable_ptr().untyped(),
+                        message: "Expected a string literal".to_string(),
+                        severity: Severity::Error,
+                    });
+                }
+            }
+        }
+        _ => {
+            return Err(PluginDiagnostic {
+                stable_ptr: expression.stable_ptr().untyped(),
+                message: "The list of namespaces should be a fixed size array of strings."
+                    .to_string(),
+                severity: Severity::Error,
+            });
+        }
+    }
+
+    Ok(namespaces)
+}

crates/dojo-lang/src/plugin.rs

@@ -30,8 +30,10 @@ use crate::event::handle_event_struct;
 use crate::inline_macros::delete::DeleteMacro;
 use crate::inline_macros::emit::EmitMacro;
 use crate::inline_macros::get::GetMacro;
+use crate::inline_macros::get_models_test_class_hashes::GetModelsTestClassHashes;
 use crate::inline_macros::selector_from_tag::SelectorFromTagMacro;
 use crate::inline_macros::set::SetMacro;
+use crate::inline_macros::spawn_test_world::SpawnTestWorld;
 use crate::interface::DojoInterface;
 use crate::introspect::{handle_introspect_enum, handle_introspect_struct};
 use crate::model::handle_model_struct;
@@ -155,7 +157,9 @@ pub fn dojo_plugin_suite() -> PluginSuite {
         .add_inline_macro_plugin::<GetMacro>()
         .add_inline_macro_plugin::<SetMacro>()
         .add_inline_macro_plugin::<EmitMacro>()
-        .add_inline_macro_plugin::<SelectorFromTagMacro>();
+        .add_inline_macro_plugin::<SelectorFromTagMacro>()
+        .add_inline_macro_plugin::<GetModelsTestClassHashes>()
+        .add_inline_macro_plugin::<SpawnTestWorld>();
 
     suite
 }

crates/dojo-lang/src/scarb_internal/mod.rs

@@ -15,7 +15,9 @@ use cairo_lang_starknet::starknet_plugin_suite;
 use cairo_lang_test_plugin::test_plugin_suite;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
 use camino::Utf8PathBuf;
-use dojo_world::config::{NamespaceConfig, DEFAULT_NAMESPACE_CFG_KEY, NAMESPACE_CFG_PREFIX};
+use dojo_world::config::{
+    NamespaceConfig, DEFAULT_NAMESPACE_CFG_KEY, DOJO_MANIFESTS_DIR_CFG_KEY, NAMESPACE_CFG_PREFIX,
+};
 use dojo_world::metadata::dojo_metadata_from_package;
 use regex::Regex;
 use scarb::compiler::{
@@ -110,6 +112,7 @@ pub fn compile_workspace(
         .collect::<Vec<_>>();
 
     let mut compile_error_units = vec![];
+
     for unit in compilation_units {
         trace!(target: LOG_TARGET, unit_name = %unit.name(), target_kind = %unit.main_component().target_kind(), "Compiling unit.");
 
@@ -208,6 +211,14 @@ pub fn cfg_set_from_component(
     let cname = c.cairo_package_name().clone();
     let package_dojo_metadata = dojo_metadata_from_package(&c.package, ws)?;
 
+    let dojo_manifests_dir = ws
+        .config()
+        .manifest_path()
+        .parent()
+        .expect("Scarb.toml manifest should always have parent")
+        .join("manifests")
+        .join(ws.current_profile().expect("profile should be set").to_string());
+
     ui.verbose(format!("component: {} ({})", cname, c.package.manifest_path()));
 
     tracing::debug!(target: LOG_TARGET, ?c, ?package_dojo_metadata);
@@ -226,6 +237,11 @@ pub fn cfg_set_from_component(
     // Add it's name for debugging on the plugin side.
     cfg_set.insert(component_cfg);
 
+    cfg_set.insert(Cfg {
+        key: DOJO_MANIFESTS_DIR_CFG_KEY.into(),
+        value: Some(dojo_manifests_dir.to_string().into()),
+    });
+
     cfg_set.insert(Cfg {
         key: DEFAULT_NAMESPACE_CFG_KEY.into(),
         value: Some(package_dojo_metadata.namespace.default.into()),