diff --git a/src/plugins_exts_yangdata.c b/src/plugins_exts_yangdata.c
index 75c4e57..ec4e8fc 100644
--- a/src/plugins_exts_yangdata.c
+++ b/src/plugins_exts_yangdata.c
@@ -58,13 +58,6 @@
                 p_ext->name, ly_stmt2str(c_ext->parent_stmt));
         return LY_ENOT;
     }
-    /* check mandatory argument */
-    if (!c_ext->argument) {
-        lyext_log(c_ext, LY_LLERR, LY_EVALID, cctx->path,
-                "Extension %s is instantiated without mandatory argument representing YANG data template name.",
-                p_ext->name);
-        return LY_EVALID;
-    }
 
     mod_c = (struct lysc_module *)c_ext->parent;
 
diff --git a/src/printer_yang.c b/src/printer_yang.c
index fbaa2e6..ccf321b 100644
--- a/src/printer_yang.c
+++ b/src/printer_yang.c
@@ -209,12 +209,13 @@
     LY_ARRAY_COUNT_TYPE u;
     struct lysp_stmt *stmt;
     ly_bool child_presence;
-    const char *argument;
 
     if (!count && ext) {
         count = LY_ARRAY_COUNT(ext);
     }
     LY_ARRAY_FOR(ext, u) {
+        struct lysp_ext *ext_def = NULL;
+
         if (!count) {
             break;
         }
@@ -224,16 +225,17 @@
             continue;
         }
 
-        ypr_open(ctx->out, flag);
-        argument = NULL;
-        if (ext[u].compiled) {
-            argument = ext[u].compiled->argument;
-        } else {
-            argument = ext[u].argument;
+        lysp_ext_find_definition(ctx->module->ctx, &ext[u], NULL, &ext_def);
+        if (!ext_def) {
+            continue;
         }
-        if (argument) {
+
+        ypr_open(ctx->out, flag);
+
+        if (ext_def->argument) {
             ly_print_(ctx->out, "%*s%s \"", INDENT, ext[u].name);
-            ypr_encode(ctx->out, argument, -1);
+            lysp_ext_instance_resolve_argument(ctx->module->ctx, &ext[u], ext_def);
+            ypr_encode(ctx->out, ext[u].argument, -1);
             ly_print_(ctx->out, "\"");
         } else {
             ly_print_(ctx->out, "%*s%s", INDENT, ext[u].name);
diff --git a/src/printer_yin.c b/src/printer_yin.c
index e1154d0..208e60a 100644
--- a/src/printer_yin.c
+++ b/src/printer_yin.c
@@ -1248,14 +1248,14 @@
 {
     LY_ARRAY_COUNT_TYPE u;
     struct lysp_stmt *stmt;
-    const char *argument;
-    const char *ext_argument;
     int8_t inner_flag = 0;
 
     if (!count && ext) {
         count = LY_ARRAY_COUNT(ext);
     }
     LY_ARRAY_FOR(ext, u) {
+        struct lysp_ext *ext_def = NULL;
+
         if (!count) {
             break;
         }
@@ -1265,21 +1265,38 @@
             continue;
         }
 
-        ypr_close_parent(ctx, flag);
-        inner_flag = 0;
-        argument = NULL;
-        ext_argument = NULL;
-
-        if (ext[u].compiled) {
-            argument = ext[u].compiled->argument;
-            ext_argument = ext[u].compiled->def->argument;
-        } else {
-            argument = ext[u].argument;
+        lysp_ext_find_definition(ctx->module->ctx, &ext[u], NULL, &ext_def);
+        if (!ext_def) {
+            continue;
         }
 
-        ypr_open(ctx, ext[u].name, ext_argument, argument, inner_flag);
+        ypr_close_parent(ctx, flag);
+        inner_flag = 0;
+
+        if (ext_def->argument) {
+            lysp_ext_instance_resolve_argument(ctx->module->ctx, &ext[u], ext_def);
+        }
+
+        ypr_open(ctx, ext[u].name, (ext_def->flags & LYS_YINELEM_TRUE) ? NULL : ext_def->argument, ext[u].argument, inner_flag);
         LEVEL++;
+        if (ext_def->flags & LYS_YINELEM_TRUE) {
+            const char *prefix, *name, *id;
+            size_t prefix_len, name_len;
+
+            ypr_close_parent(ctx, &inner_flag);
+
+            /* we need to use the same namespace as for the extension instance element */
+            id = ext[u].name;
+            ly_parse_nodeid(&id, &prefix, &prefix_len, &name, &name_len);
+            ly_print_(ctx->out, "%*s<%.*s:%s>", INDENT, prefix_len, prefix, ext_def->argument);
+            lyxml_dump_text(ctx->out, ext[u].argument, 0);
+            ly_print_(ctx->out, "</%.*s:%s>\n", prefix_len, prefix, ext_def->argument);
+        }
         LY_LIST_FOR(ext[u].child, stmt) {
+            if (stmt->flags & (LYS_YIN_ATTR | LYS_YIN_ARGUMENT)) {
+                continue;
+            }
+
             ypr_close_parent(ctx, &inner_flag);
             yprp_stmt(ctx, stmt);
         }
diff --git a/src/schema_compile.c b/src/schema_compile.c
index 5db58e0..39ba42b 100644
--- a/src/schema_compile.c
+++ b/src/schema_compile.c
@@ -92,13 +92,8 @@
 lys_compile_ext(struct lysc_ctx *ctx, struct lysp_ext_instance *ext_p, struct lysc_ext_instance *ext, void *parent,
         const struct lys_module *ext_mod)
 {
-    LY_ERR r, ret = LY_SUCCESS;
-    const char *tmp, *name, *prefix;
-    size_t pref_len, name_len;
-    LY_ARRAY_COUNT_TYPE v;
-
-    DUP_STRING(ctx->ctx, ext_p->argument, ext->argument, ret);
-    LY_CHECK_RET(ret);
+    LY_ERR ret = LY_SUCCESS;
+    struct lysp_ext *ext_def;
 
     ext->parent_stmt = ext_p->parent_stmt;
     ext->parent_stmt_index = ext_p->parent_stmt_index;
@@ -108,79 +103,15 @@
     lysc_update_path(ctx, LY_STMT_IS_NODE(ext->parent_stmt) ? ((struct lysc_node *)ext->parent)->module : NULL, "{extension}");
     lysc_update_path(ctx, NULL, ext_p->name);
 
-    /* parse the prefix */
-    tmp = ext_p->name;
-    r = ly_parse_nodeid(&tmp, &prefix, &pref_len, &name, &name_len);
-    /* it was parsed already */
-    assert(!r && !tmp[0]);
+    LY_CHECK_GOTO(ret = lysp_ext_find_definition(ctx->ctx, ext_p, &ext_mod, &ext_def), cleanup);
+    LY_CHECK_GOTO(ret = lys_compile_extension(ctx, ext_mod, ext_def, &ext->def), cleanup);
 
-    /* get module where the extension definition should be placed */
-    if (!ext_mod) {
-        ext_mod = ly_resolve_prefix(ctx->ctx, prefix, pref_len, ext_p->format, ext_p->prefix_data);
-        if (!ext_mod) {
-            LOGVAL(ctx->ctx, LYVE_REFERENCE, "Invalid prefix \"%.*s\" used for extension instance identifier.",
-                    pref_len, prefix);
-            ret = LY_EVALID;
-            goto cleanup;
-        } else if (!ext_mod->parsed->extensions) {
-            LOGVAL(ctx->ctx, LYVE_REFERENCE,
-                    "Extension instance \"%s\" refers \"%s\" module that does not contain extension definitions.",
-                    ext_p->name, ext_mod->name);
-            ret = LY_EVALID;
-            goto cleanup;
-        }
+    if (ext_def->argument) {
+        LY_CHECK_GOTO(ret = lysp_ext_instance_resolve_argument(ctx->ctx, ext_p, ext_def), cleanup);
     }
 
-    /* find the parsed extension definition there */
-    LY_ARRAY_FOR(ext_mod->parsed->extensions, v) {
-        if (!strcmp(name, ext_mod->parsed->extensions[v].name)) {
-            /* compile extension definition and assign it */
-            LY_CHECK_GOTO(ret = lys_compile_extension(ctx, ext_mod, &ext_mod->parsed->extensions[v], &ext->def), cleanup);
-            break;
-        }
-    }
-    if (!ext->def) {
-        LOGVAL(ctx->ctx, LYVE_REFERENCE, "Extension definition of extension instance \"%s\" not found.", ext_p->name);
-        ret = LY_EVALID;
-        goto cleanup;
-    }
-
-    /* unify the parsed extension from YIN and YANG sources. Without extension definition, it is not possible
-     * to get extension's argument from YIN source, so it is stored as one of the substatements. Here we have
-     * to find it, mark it with LYS_YIN_ARGUMENT and store it in the compiled structure. */
-    if ((ext_p->format == LY_PREF_XML) && ext->def->argument && !ext->argument) {
-        /* Schema was parsed from YIN and an argument is expected, ... */
-        struct lysp_stmt *stmt = NULL;
-
-        if (ext->def->flags & LYS_YINELEM_TRUE) {
-            /* ... argument was the first XML child element */
-            if (ext_p->child && !(ext_p->child->flags & LYS_YIN_ATTR)) {
-                /* TODO check namespace of the statement */
-                if (!strcmp(ext_p->child->stmt, ext->def->argument)) {
-                    stmt = ext_p->child;
-                }
-            }
-        } else {
-            /* ... argument was one of the XML attributes which are represented as child stmt
-             * with LYS_YIN_ATTR flag */
-            for (stmt = ext_p->child; stmt && (stmt->flags & LYS_YIN_ATTR); stmt = stmt->next) {
-                if (!strcmp(stmt->stmt, ext->def->argument)) {
-                    /* this is the extension's argument */
-                    break;
-                }
-            }
-        }
-        if (!stmt) {
-            /* missing extension's argument */
-            LOGVAL(ctx->ctx, LYVE_REFERENCE, "Extension instance \"%s\" misses argument \"%s\".",
-                    ext_p->name, ext->def->argument);
-            ret = LY_EVALID;
-            goto cleanup;
-
-        }
-        LY_CHECK_GOTO(ret = lydict_insert(ctx->ctx, stmt->arg, 0, &ext->argument), cleanup);
-        stmt->flags |= LYS_YIN_ARGUMENT;
-    }
+    DUP_STRING(ctx->ctx, ext_p->argument, ext->argument, ret);
+    LY_CHECK_RET(ret);
 
     if (ext->def->plugin && ext->def->plugin->compile) {
         if (ext->argument) {
@@ -195,7 +126,6 @@
         }
         LY_CHECK_GOTO(ret, cleanup);
     }
-    ext_p->compiled = ext;
 
 cleanup:
     lysc_update_path(ctx, NULL, NULL);
diff --git a/src/tree_schema.h b/src/tree_schema.h
index c4279b7..dc7772f 100644
--- a/src/tree_schema.h
+++ b/src/tree_schema.h
@@ -512,8 +512,6 @@
     void *parent;                           /**< pointer to the parent element holding the extension instance(s), use
                                                  ::lysp_ext_instance#parent_stmt to access the schema element */
     struct lysp_stmt *child;                /**< list of the extension's substatements (linked list) */
-    struct lysc_ext_instance *compiled;     /**< pointer to the compiled data if any - in case the source format is YIN,
-                                                 some of the information (argument) are available only after compilation */
     enum ly_stmt parent_stmt;               /**< value identifying placement of the extension instance */
     LY_ARRAY_COUNT_TYPE parent_stmt_index;  /**< in case the instance is in a substatement, this identifies
                                                  the index of that substatement in its [sized array](@ref sizedarrays) (if any) */
diff --git a/src/tree_schema_helpers.c b/src/tree_schema_helpers.c
index 99f97e4..2845c69 100644
--- a/src/tree_schema_helpers.c
+++ b/src/tree_schema_helpers.c
@@ -1818,6 +1818,125 @@
     return result;
 }
 
+LY_ERR
+lysp_ext_find_definition(const struct ly_ctx *ctx, const struct lysp_ext_instance *ext, const struct lys_module **ext_mod,
+        struct lysp_ext **ext_def)
+{
+    LY_ERR r;
+    const char *tmp, *name, *prefix;
+    size_t pref_len, name_len;
+    LY_ARRAY_COUNT_TYPE v;
+    const struct lys_module *mod = NULL;
+
+    *ext_def = NULL;
+    if (ext_mod) {
+        *ext_mod = NULL;
+    }
+
+    /* parse the prefix */
+    tmp = ext->name;
+    r = ly_parse_nodeid(&tmp, &prefix, &pref_len, &name, &name_len);
+    /* it was parsed already */
+    assert(!r && !tmp[0]);
+
+    /* get module where the extension definition should be placed */
+    mod = ly_resolve_prefix(ctx, prefix, pref_len, ext->format, ext->prefix_data);
+    if (!mod) {
+        LOGVAL(ctx, LYVE_REFERENCE, "Invalid prefix \"%.*s\" used for extension instance identifier.", pref_len, prefix);
+        return LY_EVALID;
+    } else if (!mod->parsed->extensions) {
+        LOGVAL(ctx, LYVE_REFERENCE, "Extension instance \"%s\" refers \"%s\" module that does not contain extension definitions.",
+                ext->name, mod->name);
+        return LY_EVALID;
+    }
+
+    /* find the parsed extension definition there */
+    LY_ARRAY_FOR(mod->parsed->extensions, v) {
+        if (!strcmp(name, mod->parsed->extensions[v].name)) {
+            *ext_def = &mod->parsed->extensions[v];
+            break;
+        }
+    }
+
+    if (!ext_def) {
+        LOGVAL(ctx, LYVE_REFERENCE, "Extension definition of extension instance \"%s\" not found.", ext->name);
+        return LY_EVALID;
+    }
+
+    if (ext_mod) {
+        *ext_mod = mod;
+    }
+    return LY_SUCCESS;
+}
+
+LY_ERR
+lysp_ext_instance_resolve_argument(struct ly_ctx *ctx, struct lysp_ext_instance *ext_p, struct lysp_ext *ext_def)
+{
+    if (!ext_def->argument || ext_p->argument) {
+        /* nothing to do */
+        return LY_SUCCESS;
+    }
+
+    if (ext_p->format == LY_PREF_XML) {
+        /* Schema was parsed from YIN and an argument is expected, ... */
+        struct lysp_stmt *stmt = NULL;
+
+        if (ext_def->flags & LYS_YINELEM_TRUE) {
+            /* ... argument was the first XML child element */
+            for (stmt = ext_p->child; stmt && (stmt->flags & LYS_YIN_ATTR); stmt = stmt->next) {}
+            if (stmt) {
+                const char *arg, *ext, *name_arg, *name_ext, *prefix_arg, *prefix_ext;
+                size_t name_arg_len, name_ext_len, prefix_arg_len, prefix_ext_len;
+
+                stmt = ext_p->child;
+
+                arg = stmt->stmt;
+                ly_parse_nodeid(&arg, &prefix_arg, &prefix_arg_len, &name_arg, &name_arg_len);
+                if (ly_strncmp(ext_def->argument, name_arg, name_arg_len)) {
+                    LOGVAL(ctx, LYVE_SEMANTICS, "Extension instance \"%s\" expects argument element \"%s\" as its first XML child, "
+                            "but \"%.*s\" element found.", ext_p->name, ext_def->argument, name_arg_len, name_arg);
+                    return LY_EVALID;
+                }
+
+                /* check namespace - all the extension instances must be qualified and argument element is expected in the same
+                 * namespace. Do not check just prefixes, there can be different prefixes pointing to the same namespace */
+                ext = ext_p->name; /* include prefix */
+                ly_parse_nodeid(&ext, &prefix_ext, &prefix_ext_len, &name_ext, &name_ext_len);
+
+                if (ly_resolve_prefix(ctx, prefix_ext, prefix_ext_len, ext_p->format, ext_p->prefix_data) !=
+                        ly_resolve_prefix(ctx, prefix_arg, prefix_arg_len, stmt->format, stmt->prefix_data)) {
+                    LOGVAL(ctx, LYVE_SEMANTICS, "Extension instance \"%s\" element and its argument element \"%s\" are "
+                            "expected in the same namespace, but they differ.", ext_p->name, ext_def->argument);
+                    return LY_EVALID;
+                }
+            }
+        } else {
+            /* ... argument was one of the XML attributes which are represented as child stmt
+             * with LYS_YIN_ATTR flag */
+            for (stmt = ext_p->child; stmt && (stmt->flags & LYS_YIN_ATTR); stmt = stmt->next) {
+                if (!strcmp(stmt->stmt, ext_def->argument)) {
+                    /* this is the extension's argument */
+                    break;
+                }
+            }
+        }
+
+        if (stmt) {
+            LY_CHECK_RET(lydict_insert(ctx, stmt->arg, 0, &ext_p->argument));
+            stmt->flags |= LYS_YIN_ARGUMENT;
+        }
+    }
+
+    if (!ext_p->argument) {
+        /* missing extension's argument */
+        LOGVAL(ctx, LYVE_SEMANTICS, "Extension instance \"%s\" misses argument %s\"%s\".",
+                ext_p->name, (ext_def->flags & LYS_YINELEM_TRUE) ? "element " : "", ext_def->argument);
+        return LY_EVALID;
+    }
+
+    return LY_SUCCESS;
+}
+
 LY_ARRAY_COUNT_TYPE
 lysp_ext_instance_iter(struct lysp_ext_instance *ext, LY_ARRAY_COUNT_TYPE index, enum ly_stmt substmt)
 {
diff --git a/src/tree_schema_internal.h b/src/tree_schema_internal.h
index 1e4a9c6..062e1d5 100644
--- a/src/tree_schema_internal.h
+++ b/src/tree_schema_internal.h
@@ -447,6 +447,34 @@
 struct lysc_must **lysc_node_musts_p(const struct lysc_node *node);
 
 /**
+ * @brief Find parsed extension definition for the given extension instance.
+ *
+ * @param[in] ctx libyang context.
+ * @param[in] ext Extension instance for which the definition will be searched.
+ * @param[in, out] ext_mod Pointer to the module where the extension definition of the @p ext to correctly resolve prefixes.
+ * @param[out] ext_def Pointer to return found extension definition.
+ * @return LY_SUCCESS when the definition was found.
+ * @return LY_EVALID when the extension instance is invalid and/or the definition not found.
+ */
+LY_ERR lysp_ext_find_definition(const struct ly_ctx *ctx, const struct lysp_ext_instance *ext, const struct lys_module **ext_mod,
+        struct lysp_ext **ext_def);
+
+/**
+ * @brief When the module comes from YIN format, the argument name is unknown because of missing extension definition
+ * (it might come from import modules which is not yet parsed at that time). Therefore, all the attributes are stored
+ * as substatements and resolving argument is postponed.
+ *
+ * There are 3 places which need the argument, so they resolve it when missing - YIN and YANG printers and extension instance
+ * compiler.
+ *
+ * @param[in] ctx libyang context
+ * @param[in] ext_p Parsed extension to be updated.
+ * @param[in] ext_def Extension definition, found with ::lysp_ext_find_definition().
+ * @return LY_ERR value.
+ */
+LY_ERR lysp_ext_instance_resolve_argument(struct ly_ctx *ctx, struct lysp_ext_instance *ext_p, struct lysp_ext *ext_def);
+
+/**
  * @brief Iterate over the specified type of the extension instances
  *
  * @param[in] ext ([Sized array](@ref sizedarrays)) of extensions to explore
diff --git a/tests/utests/CMakeLists.txt b/tests/utests/CMakeLists.txt
index 4e3543f..5b17b5c 100644
--- a/tests/utests/CMakeLists.txt
+++ b/tests/utests/CMakeLists.txt
@@ -24,7 +24,7 @@
 ly_add_utest(NAME xpath SOURCES basic/test_xpath.c)
 ly_add_utest(NAME yanglib SOURCES basic/test_yanglib.c)
 
-ly_add_utest(NAME schema SOURCES schema/test_schema.c schema/test_schema_common.c schema/test_schema_stmts.c)
+ly_add_utest(NAME schema SOURCES schema/test_schema.c schema/test_schema_common.c schema/test_schema_stmts.c schema/test_schema_extensions.c)
 ly_add_utest(NAME parser_yang SOURCES schema/test_parser_yang.c)
 ly_add_utest(NAME parser_yin SOURCES schema/test_parser_yin.c)
 ly_add_utest(NAME tree_schema_compile SOURCES schema/test_tree_schema_compile.c)
diff --git a/tests/utests/extensions/test_metadata.c b/tests/utests/extensions/test_metadata.c
index ca5558c..163785b 100644
--- a/tests/utests/extensions/test_metadata.c
+++ b/tests/utests/extensions/test_metadata.c
@@ -125,7 +125,7 @@
             "<md:annotation name=\"aa\"/>\n"
             "</module>";
     assert_int_equal(LY_EVALID, lys_parse_mem(UTEST_LYCTX, data, LYS_IN_YIN, NULL));
-    CHECK_LOG_CTX("Missing mandatory keyword \"type\" as a child of \"md:annotation\".", "/aa:{extension='md:annotation'}/aa");
+    CHECK_LOG_CTX("Missing mandatory keyword \"type\" as a child of \"md:annotation aa\".", "/aa:{extension='md:annotation'}/aa");
 
     /* not allowed substatement */
     data = "<module xmlns=\"urn:ietf:params:xml:ns:yang:yin:1\" xmlns:md=\"urn:ietf:params:xml:ns:yang:ietf-yang-metadata\" name=\"aa\">\n"
@@ -135,7 +135,7 @@
             "  <default value=\"x\"/>\n"
             "</md:annotation></module>";
     assert_int_equal(LY_EVALID, lys_parse_mem(UTEST_LYCTX, data, LYS_IN_YIN, NULL));
-    CHECK_LOG_CTX("Invalid keyword \"default\" as a child of \"md:annotation\" extension instance.", "/aa:{extension='md:annotation'}/aa");
+    CHECK_LOG_CTX("Invalid keyword \"default\" as a child of \"md:annotation aa\" extension instance.", "/aa:{extension='md:annotation'}/aa");
 
     /* invalid cardinality of units substatement */
     data = "<module xmlns=\"urn:ietf:params:xml:ns:yang:yin:1\" xmlns:md=\"urn:ietf:params:xml:ns:yang:ietf-yang-metadata\" name=\"aa\">\n"
diff --git a/tests/utests/extensions/test_yangdata.c b/tests/utests/extensions/test_yangdata.c
index 02b50f1..7bf1cc6 100644
--- a/tests/utests/extensions/test_yangdata.c
+++ b/tests/utests/extensions/test_yangdata.c
@@ -210,8 +210,7 @@
             "import ietf-restconf {revision-date 2017-01-26; prefix rc;}"
             "rc:yang-data { container x { leaf x {type string;}}}}";
     assert_int_equal(LY_EVALID, lys_parse_mem(UTEST_LYCTX, data, LYS_IN_YANG, NULL));
-    CHECK_LOG_CTX("Extension plugin \"libyang 2 - yang-data, version 1\": "
-            "Extension rc:yang-data is instantiated without mandatory argument representing YANG data template name.",
+    CHECK_LOG_CTX("Extension instance \"rc:yang-data\" misses argument element \"name\".",
             "/a:{extension='rc:yang-data'}");
 
     data = "module a {yang-version 1.1; namespace urn:tests:extensions:yangdata:a; prefix self;"
diff --git a/tests/utests/schema/test_parser_yin.c b/tests/utests/schema/test_parser_yin.c
index 14f6b7d..9292f0d 100644
--- a/tests/utests/schema/test_parser_yin.c
+++ b/tests/utests/schema/test_parser_yin.c
@@ -121,7 +121,7 @@
 #define ELEMENT_WRAPPER_END "</status>"
 
 #define TEST_1_CHECK_LYSP_EXT_INSTANCE(NODE, INSUBSTMT)\
-    CHECK_LYSP_EXT_INSTANCE((NODE), NULL, 1, NULL, INSUBSTMT, 0, "myext:c-define", LY_PREF_XML)
+    CHECK_LYSP_EXT_INSTANCE((NODE), NULL, 1, INSUBSTMT, 0, "myext:c-define", LY_PREF_XML)
 
 struct lys_yin_parser_ctx *YCTX;
 
@@ -335,7 +335,7 @@
 
     ret = yin_parse_extension_instance(YCTX, LY_STMT_CONTACT, 0, &exts);
     assert_int_equal(ret, LY_SUCCESS);
-    CHECK_LYSP_EXT_INSTANCE(exts, NULL, 1, NULL, LY_STMT_CONTACT, 0, "myext:ext", LY_PREF_XML);
+    CHECK_LYSP_EXT_INSTANCE(exts, NULL, 1, LY_STMT_CONTACT, 0, "myext:ext", LY_PREF_XML);
 
     CHECK_LYSP_STMT(exts->child, arg, 0, LYS_YIN_ATTR, 0, 1, stmt);
     stmt = "value";
@@ -355,7 +355,7 @@
 
     ret = yin_parse_extension_instance(YCTX, LY_STMT_CONTACT, 0, &exts);
     assert_int_equal(ret, LY_SUCCESS);
-    CHECK_LYSP_EXT_INSTANCE(exts, NULL, 0, NULL, LY_STMT_CONTACT, 0, "myext:extension-elem", LY_PREF_XML);
+    CHECK_LYSP_EXT_INSTANCE(exts, NULL, 0, LY_STMT_CONTACT, 0, "myext:extension-elem", LY_PREF_XML);
     lysp_ext_instance_free(UTEST_LYCTX, exts);
     LY_ARRAY_FREE(exts);
     exts = NULL;
@@ -377,7 +377,7 @@
     ret = yin_parse_extension_instance(YCTX, LY_STMT_CONTACT, 0, &exts);
     assert_int_equal(ret, LY_SUCCESS);
 
-    CHECK_LYSP_EXT_INSTANCE(exts, NULL, 1, NULL, LY_STMT_CONTACT, 0, "myext:ext", LY_PREF_XML);
+    CHECK_LYSP_EXT_INSTANCE(exts, NULL, 1, LY_STMT_CONTACT, 0, "myext:ext", LY_PREF_XML);
 
     stmt = "attr1";
     arg = "text1";
@@ -544,7 +544,7 @@
     const char *exts_name = "myext:custom";
     const char *exts_arg = "totally amazing extension";
 
-    CHECK_LYSP_EXT_INSTANCE(exts, exts_arg, 0, NULL, LY_STMT_PREFIX, 0, exts_name, LY_PREF_XML);
+    CHECK_LYSP_EXT_INSTANCE(exts, exts_arg, 0, LY_STMT_PREFIX, 0, exts_name, LY_PREF_XML);
     assert_string_equal(value, "wsefsdf");
     assert_string_equal(units, "radians");
     assert_string_equal(when_p->cond, "condition...");
diff --git a/tests/utests/schema/test_schema.c b/tests/utests/schema/test_schema.c
index dd9e34c..3c49977 100644
--- a/tests/utests/schema/test_schema.c
+++ b/tests/utests/schema/test_schema.c
@@ -26,7 +26,11 @@
         const char **module_data, void (**free_module_data)(void *model_data, void *user_data))
 {
     *module_data = user_data;
-    *format = LYS_IN_YANG;
+    if ((*module_data)[0] == '<') {
+        *format = LYS_IN_YIN;
+    } else {
+        *format = LYS_IN_YANG;
+    }
     *free_module_data = NULL;
     return LY_SUCCESS;
 }
@@ -46,6 +50,10 @@
 void test_identity(void **state);
 void test_feature(void **state);
 
+/* test_schema_extensions.c */
+void test_extension_argument(void **state);
+void test_extension_argument_element(void **state);
+
 int
 main(void)
 {
@@ -61,6 +69,10 @@
         /** test_schema_stmts.c */
         UTEST(test_identity),
         UTEST(test_feature),
+
+        /** test_schema_extensions.c */
+        UTEST(test_extension_argument),
+        UTEST(test_extension_argument_element),
     };
 
     return cmocka_run_group_tests(tests, NULL, NULL);
diff --git a/tests/utests/schema/test_schema_extensions.c b/tests/utests/schema/test_schema_extensions.c
new file mode 100644
index 0000000..7c15773
--- /dev/null
+++ b/tests/utests/schema/test_schema_extensions.c
@@ -0,0 +1,288 @@
+/*
+ * @file test_schema_extensions.c
+ * @author: Radek Krejci <rkrejci@cesnet.cz>
+ * @brief unit tests for YANG (YIN) extension statements and their instances in schemas
+ *
+ * Copyright (c) 2018-2021 CESNET, z.s.p.o.
+ *
+ * This source code is licensed under BSD 3-Clause License (the "License").
+ * You may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://opensource.org/licenses/BSD-3-Clause
+ */
+#include "test_schema.h"
+
+#include <string.h>
+
+#include "context.h"
+#include "log.h"
+#include "tree_schema.h"
+
+void
+test_extension_argument(void **state)
+{
+    const struct lys_module *mod;
+    const char *mod_def_yang = "module a {\n"
+            "  namespace \"urn:a\";\n"
+            "  prefix a;\n\n"
+            "  extension e {\n"
+            "    argument name;\n"
+            "  }\n\n"
+            "  a:e \"aaa\";\n"
+            "}\n";
+    const char *mod_def_yin =
+            "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"
+            "<module name=\"a\"\n"
+            "        xmlns=\"urn:ietf:params:xml:ns:yang:yin:1\"\n"
+            "        xmlns:a=\"urn:a\">\n"
+            "  <namespace uri=\"urn:a\"/>\n"
+            "  <prefix value=\"a\"/>\n\n"
+            "  <extension name=\"e\">\n"
+            "    <argument name=\"name\"/>\n"
+            "  </extension>\n\n"
+            "  <a:e name=\"aaa\"/>\n"
+            "</module>\n";
+    const char *mod_test_yin, *mod_test_yang;
+    char *printed;
+
+    mod_test_yang = "module b {\n"
+            "  namespace \"urn:b\";\n"
+            "  prefix b;\n\n"
+            "  import a {\n"
+            "    prefix a;\n"
+            "  }\n\n"
+            "  a:e \"xxx\";\n"
+            "}\n";
+    mod_test_yin =
+            "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"
+            "<module name=\"b\"\n"
+            "        xmlns=\"urn:ietf:params:xml:ns:yang:yin:1\"\n"
+            "        xmlns:b=\"urn:b\"\n"
+            "        xmlns:a=\"urn:a\">\n"
+            "  <namespace uri=\"urn:b\"/>\n"
+            "  <prefix value=\"b\"/>\n"
+            "  <import module=\"a\">\n"
+            "    <prefix value=\"a\"/>\n"
+            "  </import>\n\n"
+            "  <a:e name=\"xxx\"/>\n"
+            "</module>\n";
+
+    /* from YANG */
+    ly_ctx_set_module_imp_clb(UTEST_LYCTX, test_imp_clb, (void *)mod_def_yang);
+    assert_int_equal(LY_SUCCESS, lys_parse_mem(UTEST_LYCTX, mod_test_yang, LYS_IN_YANG, &mod));
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YANG, 0));
+    assert_string_equal(printed, mod_test_yang);
+    free(printed);
+
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YIN, 0));
+    assert_string_equal(printed, mod_test_yin);
+    free(printed);
+
+    assert_non_null(mod = ly_ctx_get_module(UTEST_LYCTX, "a", NULL));
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YANG, 0));
+    assert_string_equal(printed, mod_def_yang);
+    free(printed);
+
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YIN, 0));
+    assert_string_equal(printed, mod_def_yin);
+    free(printed);
+
+    /* context reset */
+    ly_ctx_destroy(UTEST_LYCTX, NULL);
+    ly_ctx_new(NULL, 0, &UTEST_LYCTX);
+
+    /* from YIN */
+    ly_ctx_set_module_imp_clb(UTEST_LYCTX, test_imp_clb, (void *)mod_def_yin);
+    assert_int_equal(LY_SUCCESS, lys_parse_mem(UTEST_LYCTX, mod_test_yin, LYS_IN_YIN, &mod));
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YANG, 0));
+    assert_string_equal(printed, mod_test_yang);
+    free(printed);
+
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YIN, 0));
+    assert_string_equal(printed, mod_test_yin);
+    free(printed);
+
+    assert_non_null(mod = ly_ctx_get_module(UTEST_LYCTX, "a", NULL));
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YANG, 0));
+    assert_string_equal(printed, mod_def_yang);
+    free(printed);
+
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YIN, 0));
+    assert_string_equal(printed, mod_def_yin);
+    free(printed);
+}
+
+void
+test_extension_argument_element(void **state)
+{
+    const struct lys_module *mod;
+    const char *mod_def_yang = "module a {\n"
+            "  namespace \"urn:a\";\n"
+            "  prefix a;\n\n"
+            "  extension e {\n"
+            "    argument name {\n"
+            "      yin-element true;\n"
+            "    }\n"
+            "  }\n\n"
+            "  a:e \"aaa\";\n"
+            "}\n";
+    const char *mod_def_yin =
+            "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"
+            "<module name=\"a\"\n"
+            "        xmlns=\"urn:ietf:params:xml:ns:yang:yin:1\"\n"
+            "        xmlns:a=\"urn:a\">\n"
+            "  <namespace uri=\"urn:a\"/>\n"
+            "  <prefix value=\"a\"/>\n\n"
+            "  <extension name=\"e\">\n"
+            "    <argument name=\"name\">\n"
+            "      <yin-element value=\"true\"/>\n"
+            "    </argument>\n"
+            "  </extension>\n\n"
+            "  <a:e>\n"
+            "    <a:name>aaa</a:name>\n"
+            "  </a:e>\n"
+            "</module>\n";
+    const char *mod_test_yin, *mod_test_yang;
+    char *printed;
+
+    mod_test_yang = "module b {\n"
+            "  namespace \"urn:b\";\n"
+            "  prefix b;\n\n"
+            "  import a {\n"
+            "    prefix a;\n"
+            "  }\n\n"
+            "  a:e \"xxx\";\n"
+            "}\n";
+    mod_test_yin =
+            "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"
+            "<module name=\"b\"\n"
+            "        xmlns=\"urn:ietf:params:xml:ns:yang:yin:1\"\n"
+            "        xmlns:b=\"urn:b\"\n"
+            "        xmlns:a=\"urn:a\">\n"
+            "  <namespace uri=\"urn:b\"/>\n"
+            "  <prefix value=\"b\"/>\n"
+            "  <import module=\"a\">\n"
+            "    <prefix value=\"a\"/>\n"
+            "  </import>\n\n"
+            "  <a:e>\n"
+            "    <a:name>xxx</a:name>\n"
+            "  </a:e>\n"
+            "</module>\n";
+
+    /* from YANG */
+    ly_ctx_set_module_imp_clb(UTEST_LYCTX, test_imp_clb, (void *)mod_def_yang);
+    assert_int_equal(LY_SUCCESS, lys_parse_mem(UTEST_LYCTX, mod_test_yang, LYS_IN_YANG, &mod));
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YANG, 0));
+    assert_string_equal(printed, mod_test_yang);
+    free(printed);
+
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YIN, 0));
+    assert_string_equal(printed, mod_test_yin);
+    free(printed);
+
+    assert_non_null(mod = ly_ctx_get_module(UTEST_LYCTX, "a", NULL));
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YANG, 0));
+    assert_string_equal(printed, mod_def_yang);
+    free(printed);
+
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YIN, 0));
+    assert_string_equal(printed, mod_def_yin);
+    free(printed);
+
+    /* context reset */
+    ly_ctx_destroy(UTEST_LYCTX, NULL);
+    ly_ctx_new(NULL, 0, &UTEST_LYCTX);
+
+    /* from YIN */
+    ly_ctx_set_module_imp_clb(UTEST_LYCTX, test_imp_clb, (void *)mod_def_yin);
+    assert_int_equal(LY_SUCCESS, lys_parse_mem(UTEST_LYCTX, mod_test_yin, LYS_IN_YIN, &mod));
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YANG, 0));
+    assert_string_equal(printed, mod_test_yang);
+    free(printed);
+
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YIN, 0));
+    assert_string_equal(printed, mod_test_yin);
+    free(printed);
+
+    assert_non_null(mod = ly_ctx_get_module(UTEST_LYCTX, "a", NULL));
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YANG, 0));
+    assert_string_equal(printed, mod_def_yang);
+    free(printed);
+
+    assert_int_equal(LY_SUCCESS, lys_print_mem(&printed, mod, LYS_OUT_YIN, 0));
+    assert_string_equal(printed, mod_def_yin);
+    free(printed);
+
+    /* invalid */
+    mod_test_yang = "module x { namespace \"urn:x\"; prefix x; import a { prefix a; } a:e; }";
+    assert_int_equal(LY_EVALID, lys_parse_mem(UTEST_LYCTX, mod_test_yang, LYS_IN_YANG, NULL));
+    CHECK_LOG_CTX("Extension instance \"a:e\" misses argument element \"name\".", "/x:{extension='a:e'}");
+
+    mod_test_yin = "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"
+            "<module name=\"x\"\n"
+            "        xmlns=\"urn:ietf:params:xml:ns:yang:yin:1\"\n"
+            "        xmlns:x=\"urn:x\"\n"
+            "        xmlns:a=\"urn:a\">\n"
+            "  <namespace uri=\"urn:x\"/>\n"
+            "  <prefix value=\"x\"/>\n"
+            "  <import module=\"a\">\n"
+            "    <prefix value=\"a\"/>\n"
+            "  </import>\n\n"
+            "  <a:e/>\n"
+            "</module>\n";
+    assert_int_equal(LY_EVALID, lys_parse_mem(UTEST_LYCTX, mod_test_yin, LYS_IN_YIN, NULL));
+    CHECK_LOG_CTX("Extension instance \"a:e\" misses argument element \"name\".", "/x:{extension='a:e'}");
+
+    mod_test_yin = "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"
+            "<module name=\"x\"\n"
+            "        xmlns=\"urn:ietf:params:xml:ns:yang:yin:1\"\n"
+            "        xmlns:x=\"urn:x\"\n"
+            "        xmlns:a=\"urn:a\">\n"
+            "  <namespace uri=\"urn:x\"/>\n"
+            "  <prefix value=\"x\"/>\n"
+            "  <import module=\"a\">\n"
+            "    <prefix value=\"a\"/>\n"
+            "  </import>\n\n"
+            "  <a:e name=\"xxx\"/>\n"
+            "</module>\n";
+    assert_int_equal(LY_EVALID, lys_parse_mem(UTEST_LYCTX, mod_test_yin, LYS_IN_YIN, NULL));
+    CHECK_LOG_CTX("Extension instance \"a:e\" misses argument element \"name\".", "/x:{extension='a:e'}");
+
+    mod_test_yin = "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"
+            "<module name=\"x\"\n"
+            "        xmlns=\"urn:ietf:params:xml:ns:yang:yin:1\"\n"
+            "        xmlns:x=\"urn:x\"\n"
+            "        xmlns:a=\"urn:a\">\n"
+            "  <namespace uri=\"urn:x\"/>\n"
+            "  <prefix value=\"x\"/>\n"
+            "  <import module=\"a\">\n"
+            "    <prefix value=\"a\"/>\n"
+            "  </import>\n\n"
+            "  <a:e>\n"
+            "    <x:name>xxx</x:name>\n"
+            "  </a:e>\n"
+            "</module>\n";
+    assert_int_equal(LY_EVALID, lys_parse_mem(UTEST_LYCTX, mod_test_yin, LYS_IN_YIN, NULL));
+    CHECK_LOG_CTX("Extension instance \"a:e\" element and its argument element \"name\" are expected in the same namespace, but they differ.",
+            "/x:{extension='a:e'}");
+
+    mod_test_yin = "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"
+            "<module name=\"x\"\n"
+            "        xmlns=\"urn:ietf:params:xml:ns:yang:yin:1\"\n"
+            "        xmlns:x=\"urn:x\"\n"
+            "        xmlns:a=\"urn:a\">\n"
+            "  <namespace uri=\"urn:x\"/>\n"
+            "  <prefix value=\"x\"/>\n"
+            "  <import module=\"a\">\n"
+            "    <prefix value=\"a\"/>\n"
+            "  </import>\n\n"
+            "  <a:e>\n"
+            "    <a:value>xxx</a:value>\n"
+            "  </a:e>\n"
+            "</module>\n";
+    assert_int_equal(LY_EVALID, lys_parse_mem(UTEST_LYCTX, mod_test_yin, LYS_IN_YIN, NULL));
+    CHECK_LOG_CTX("Extension instance \"a:e\" expects argument element \"name\" as its first XML child, but \"value\" element found.",
+            "/x:{extension='a:e'}");
+
+}
diff --git a/tests/utests/utests.h b/tests/utests/utests.h
index b57df7b..1ed4fb4 100644
--- a/tests/utests/utests.h
+++ b/tests/utests/utests.h
@@ -427,16 +427,14 @@
  * @param[in] NODE      pointer to lysp_ext_instance variable
  * @param[in] ARGUMENT  expected optional value of the extension's argument
  * @param[in] CHILD     0 -> node doesnt have child, 1 -> node have children
- * @param[in] COMPILED  0 -> compiled data dosnt exists, 1 -> compiled data exists
  * @param[in] PARENT_STMT expected value identifying placement of the extension instance
  * @param[in] PARENT_STMT_INDEX expected indentifi index
  * @param[in] FORMAT    expected format
  */
-#define CHECK_LYSP_EXT_INSTANCE(NODE, ARGUMENT, CHILD, COMPILED, PARENT_STMT, PARENT_STMT_INDEX, NAME, FORMAT) \
+#define CHECK_LYSP_EXT_INSTANCE(NODE, ARGUMENT, CHILD, PARENT_STMT, PARENT_STMT_INDEX, NAME, FORMAT) \
     assert_non_null(NODE); \
     CHECK_STRING((NODE)->argument, ARGUMENT); \
     CHECK_POINTER((NODE)->child, CHILD); \
-    CHECK_POINTER((NODE)->compiled, COMPILED); \
     /*assert_int_equal((NODE)->flags, LYS_INTERNAL);*/ \
     assert_int_equal((NODE)->parent_stmt, PARENT_STMT); \
     assert_int_equal((NODE)->parent_stmt_index, PARENT_STMT_INDEX); \