Test Only

The #[test_only] attribute allows creating a module member that will only be available in the test environment. This is useful for creating initializers and other helper functions that are needed for testing but would be dangerous to expose in the production environment. The attribute can also be used to define test-only types, friends and dependencies.

module examples::registry {
    use std::vector;
    use std::type_name::{Self, TypeName};
    use sui::tx_context::TxContext;
    use sui::object::{Self, UID};

    /// Singleton object that can be created only once.
    struct TypeRegistry has key {
        id: UID,
        /// Simple counter to keep track of the number of records.
        records: vector<TypeName>,
    }

    /// Create and share a new `Registry` object.
    fun init(ctx: &mut TxContext) {
        sui::transfer::share_object(TypeRegistry {
            id: object::new(ctx),
            records: vector[]
        })
    }

    /// Add a new record to the `Registry` object with the type name of `T`.
    public fun register_type<T>(r: &mut TypeRegistry) {
        vector::push_back(&mut r.records, type_name::get<T>())
    }

    /// Getter for the `records` field.
    public fun records(registry: &TypeRegistry): &vector<TypeName> {
        &registry.records
    }

    // === Test and test-only members ===

    #[test_only]
    // The #[test_only] attribute makes the function only available in test mode.
    // Meaning that it can be used in functions marked with #[test] and in other
    // functions marked with #[test_only].
    //
    // We allow creating a TypeRegistry object for tests! Also, in tests it's often
    // a good idea to allow for customizing the object's state.
    public fun new_for_testing(ctx: &mut TxContext): TypeRegistry {
        TypeRegistry {
            id: object::new(ctx),
            records: vector[]
        }
    }

    #[test_only]
    // It's also a good idea to mark test_only functions with suffixes like
    // `_for_testing` to make it clear that they are only available in tests.
    public fun reset_for_testing(reg: &mut TypeRegistry) {
        reg.records = vector[];
    }


    // Test_only can be applied to any module member. Including structs!
    #[test_only] struct TestOnlyStruct {}

    // Dependencies can also be marked as test_only.
    #[test_only] use sui::tx_context;

    // Some dependencies are only available in test mode such as `sui::test_utils`.
    #[test_only] use sui::test_utils;

    #[test]
    // In tests we can use the `new_for_testing` function to create a new
    // `Registry` object.
    fun test_registry() {
        let ctx = &mut tx_context::dummy();
        let registry = new_for_testing(ctx);

        assert!(vector::length(records(&registry)) == 0, 0);

        // we can use `TestOnlyStruct` because it's marked with `#[test_only]`
        register_type<TestOnlyStruct>(&mut registry);

        assert!(vector::length(records(&registry)) == 1, 0);

        // super helpful utility which is only available in tests!
        test_utils::destroy(registry);
    }
}

#[test_only]
// #[test_only] can also be applied to modules. All of the module's members
// will be test_only and don't need to be marked individually.
module examples::registry_tests {
    use std::vector;
    use sui::tx_context;
    use sui::test_utils;
    use examples::registry::{Self, TestOnlyStruct};

    struct AnotherTestOnlyStruct {}

    #[test]
    fun test_with_two_structs() {
        let ctx = &mut tx_context::dummy();
        let reg = registry::new_for_testing(ctx);

        registry::register_type<TestOnlyStruct>(&mut reg);
        registry::reset_for_testing(&mut reg);

        assert!(vector::length(registry::records(&reg)) == 0, 0);

        registry::register_type<AnotherTestOnlyStruct>(&mut reg);

        assert!(vector::length(registry::records(&reg)) == 1, 0);

        test_utils::destroy(reg);
    }
}