schema BUGFIX processing extension instance argument

Due to differences in YANG and YIN, the extension instance's argument is
processed differently. So far, the argument was missing when the source
format was YIN and extension was not compiled. This was breaking
printers and limiting even other use cases of the argument.

This patch solves the problem by getting the extension definition when
needed and finishing the argument processing not finished by YIN parser.
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); \