diff --git a/lib/northbound.c b/lib/northbound.c
index 18500a8bd2..3fc1f1c2e6 100644
--- a/lib/northbound.c
+++ b/lib/northbound.c
@@ -122,8 +122,8 @@ static int nb_node_new_cb(const struct lys_node *snode, void *arg)
 	if (CHECK_FLAG(snode->nodetype, LYS_CONTAINER | LYS_LIST)) {
 		bool config_only = true;
 
-		yang_snodes_iterate_subtree(snode, nb_node_check_config_only,
-					    YANG_ITER_ALLOW_AUGMENTATIONS,
+		yang_snodes_iterate_subtree(snode, NULL,
+					    nb_node_check_config_only, 0,
 					    &config_only);
 		if (config_only)
 			SET_FLAG(nb_node->flags, F_NB_NODE_CONFIG_ONLY);
@@ -141,6 +141,7 @@ static int nb_node_new_cb(const struct lys_node *snode, void *arg)
 	 * another.
 	 */
 	nb_node->snode = snode;
+	assert(snode->priv == NULL);
 	lys_set_private(snode, nb_node);
 
 	return YANG_ITER_CONTINUE;
diff --git a/lib/yang.c b/lib/yang.c
index 9bfdcb858c..5bf7758e18 100644
--- a/lib/yang.c
+++ b/lib/yang.c
@@ -147,11 +147,15 @@ struct yang_module *yang_module_find(const char *module_name)
 }
 
 int yang_snodes_iterate_subtree(const struct lys_node *snode,
+				const struct lys_module *module,
 				yang_iterate_cb cb, uint16_t flags, void *arg)
 {
 	struct lys_node *child;
 	int ret = YANG_ITER_CONTINUE;
 
+	if (module && snode->module != module)
+		goto next;
+
 	if (CHECK_FLAG(flags, YANG_ITER_FILTER_IMPLICIT)) {
 		switch (snode->nodetype) {
 		case LYS_CASE:
@@ -214,11 +218,8 @@ next:
 		return YANG_ITER_CONTINUE;
 
 	LY_TREE_FOR (snode->child, child) {
-		if (!CHECK_FLAG(flags, YANG_ITER_ALLOW_AUGMENTATIONS)
-		    && child->parent != snode)
-			continue;
-
-		ret = yang_snodes_iterate_subtree(child, cb, flags, arg);
+		ret = yang_snodes_iterate_subtree(child, module, cb, flags,
+						  arg);
 		if (ret == YANG_ITER_STOP)
 			return ret;
 	}
@@ -233,15 +234,16 @@ int yang_snodes_iterate_module(const struct lys_module *module,
 	int ret = YANG_ITER_CONTINUE;
 
 	LY_TREE_FOR (module->data, snode) {
-		ret = yang_snodes_iterate_subtree(snode, cb, flags, arg);
+		ret = yang_snodes_iterate_subtree(snode, module, cb, flags,
+						  arg);
 		if (ret == YANG_ITER_STOP)
 			return ret;
 	}
 
 	for (uint8_t i = 0; i < module->augment_size; i++) {
 		ret = yang_snodes_iterate_subtree(
-			(const struct lys_node *)&module->augment[i], cb, flags,
-			arg);
+			(const struct lys_node *)&module->augment[i], module,
+			cb, flags, arg);
 		if (ret == YANG_ITER_STOP)
 			return ret;
 	}
@@ -255,9 +257,14 @@ int yang_snodes_iterate_all(yang_iterate_cb cb, uint16_t flags, void *arg)
 	int ret = YANG_ITER_CONTINUE;
 
 	RB_FOREACH (module, yang_modules, &yang_modules) {
-		ret = yang_snodes_iterate_module(module->info, cb, flags, arg);
-		if (ret == YANG_ITER_STOP)
-			return ret;
+		struct lys_node *snode;
+
+		LY_TREE_FOR (module->info->data, snode) {
+			ret = yang_snodes_iterate_subtree(snode, NULL, cb,
+							  flags, arg);
+			if (ret == YANG_ITER_STOP)
+				return ret;
+		}
 	}
 
 	return ret;
diff --git a/lib/yang.h b/lib/yang.h
index 94bbed233d..867ade9676 100644
--- a/lib/yang.h
+++ b/lib/yang.h
@@ -102,9 +102,6 @@ enum yang_iter_flags {
 
 	/* Filter implicitely created nodes. */
 	YANG_ITER_FILTER_IMPLICIT = (1<<3),
-
-	/* Allow iteration over augmentations. */
-	YANG_ITER_ALLOW_AUGMENTATIONS = (1<<4),
 };
 
 /* Callback used by the yang_snodes_iterate_*() family of functions. */
@@ -168,6 +165,9 @@ extern void yang_module_embed(struct yang_module_embed *embed);
  * snode
  *    YANG schema node to operate on.
  *
+ * module
+ *    When set, iterate over all nodes of the specified module only.
+ *
  * cb
  *    Function to call with each schema node.
  *
@@ -181,6 +181,7 @@ extern void yang_module_embed(struct yang_module_embed *embed);
  *    The return value of the last called callback.
  */
 extern int yang_snodes_iterate_subtree(const struct lys_node *snode,
+				       const struct lys_module *module,
 				       yang_iterate_cb cb, uint16_t flags,
 				       void *arg);