Merge "tests: small deduplication"
diff --git a/src/ast_handlers.cpp b/src/ast_handlers.cpp
index 24866b2..f160d8c 100644
--- a/src/ast_handlers.cpp
+++ b/src/ast_handlers.cpp
@@ -23,3 +23,8 @@
     }
     return leafDataToString(value);
 }
+
+boost::optional<std::string> optModuleToOptString(const boost::optional<module_> module)
+{
+    return module.flat_map([] (const auto& module) { return boost::optional<std::string>(module.m_name); });
+}
diff --git a/src/ast_handlers.hpp b/src/ast_handlers.hpp
index 5b5a152..b9d7a9d 100644
--- a/src/ast_handlers.hpp
+++ b/src/ast_handlers.hpp
@@ -44,38 +44,23 @@
     }
 };
 
-struct node_identifier_class {
-    template <typename T, typename Iterator, typename Context>
-    void on_success(Iterator const&, Iterator const&, T&, Context const& context)
-    {
-        auto& parserContext = x3::get<parser_context_tag>(context);
+struct node_identifier_class;
 
-        if (!parserContext.m_topLevelModulePresent) {
-            if (parserContext.m_errorMsg.empty())
-                parserContext.m_errorMsg = "You have to specify a top level module.";
-            _pass(context) = false;
-        }
-    }
-};
+boost::optional<std::string> optModuleToOptString(const boost::optional<module_> module);
 
 struct key_identifier_class {
     template <typename T, typename Iterator, typename Context>
     void on_success(Iterator const&, Iterator const&, T& ast, Context const& context)
     {
-        auto& parserContext = x3::get<parser_context_tag>(context);
+        ParserContext& parserContext = x3::get<parser_context_tag>(context);
         const Schema& schema = parserContext.m_schema;
-        schemaPath_ location = parserContext.currentSchemaPath();
-        ModuleNodePair list{parserContext.m_curModule, parserContext.m_tmpListName};
 
-        if (schema.listHasKey(location, list, ast)) {
-            schemaNode_ listNode;
-            listNode.m_prefix = parserContext.m_curModule.flat_map([] (auto mod) { return boost::optional<module_>{{mod}}; });;
-            listNode.m_suffix = list_{parserContext.m_tmpListName};
-            location.m_nodes.push_back(listNode);
-            parserContext.m_tmpListKeyLeafPath.m_location = location;
-            parserContext.m_tmpListKeyLeafPath.m_node = { parserContext.m_curModule, ast };
+        if (schema.listHasKey(dataPathToSchemaPath(parserContext.m_tmpListPath), ast)) {
+            parserContext.m_tmpListKeyLeafPath.m_location = dataPathToSchemaPath(parserContext.m_tmpListPath);
+            parserContext.m_tmpListKeyLeafPath.m_node = { optModuleToOptString(parserContext.m_tmpListPath.m_nodes.back().m_prefix), ast };
         } else {
-            parserContext.m_errorMsg = parserContext.m_tmpListName + " is not indexed by \"" + ast + "\".";
+            auto listName = std::get<list_>(parserContext.m_tmpListPath.m_nodes.back().m_suffix).m_name;
+            parserContext.m_errorMsg = listName + " is not indexed by \"" + ast + "\".";
             _pass(context) = false;
         }
     }
@@ -90,13 +75,14 @@
         auto& parserContext = x3::get<parser_context_tag>(context);
         const Schema& schema = parserContext.m_schema;
 
-        const auto& keysNeeded = schema.listKeys(parserContext.currentSchemaPath(), {parserContext.m_curModule, parserContext.m_tmpListName});
+        const auto& keysNeeded = schema.listKeys(dataPathToSchemaPath(parserContext.m_tmpListPath));
         std::set<std::string> keysSupplied;
         for (const auto& it : ast)
             keysSupplied.insert(it.first);
 
         if (keysNeeded != keysSupplied) {
-            parserContext.m_errorMsg = "Not enough keys for " + parserContext.m_tmpListName + ". " +
+            auto listName = std::get<list_>(parserContext.m_tmpListPath.m_nodes.back().m_suffix).m_name;
+            parserContext.m_errorMsg = "Not enough keys for " + listName + ". " +
                                        "These keys were not supplied:";
             std::set<std::string> missingKeys;
             std::set_difference(keysNeeded.begin(), keysNeeded.end(),
@@ -128,10 +114,7 @@
         auto& parserContext = x3::get<parser_context_tag>(context);
         const auto& schema = parserContext.m_schema;
 
-        if (schema.isModule(ast.m_name)) {
-            parserContext.m_curModule = ast.m_name;
-            parserContext.m_topLevelModulePresent = true;
-        } else {
+        if (!schema.isModule(ast.m_name)) {
             parserContext.m_errorMsg = "Invalid module name.";
             _pass(context) = false;
         }
@@ -275,22 +258,6 @@
     }
 };
 
-// This handler only checks if the module exists
-// It doesn't set any ParserContext flags (except the error message)
-struct data_module_prefix_class {
-    template <typename T, typename Iterator, typename Context>
-    void on_success(Iterator const&, Iterator const&, T& ast, Context const& context)
-    {
-        auto& parserContext = x3::get<parser_context_tag>(context);
-        const auto& schema = parserContext.m_schema;
-
-        if (!schema.isModule(parserContext.currentSchemaPath(), ast.m_name)) {
-            parserContext.m_errorMsg = "Invalid module name.";
-            _pass(context) = false;
-        }
-    }
-};
-
 struct set_class {
     template <typename Iterator, typename Exception, typename Context>
     x3::error_handler_result on_error(Iterator&, Iterator const&, Exception const& x, Context const& context)
@@ -333,7 +300,6 @@
         auto& parserContext = x3::get<parser_context_tag>(context);
         parserContext.resetPath();
         parserContext.m_tmpListKeys.clear();
-        parserContext.m_tmpListName.clear();
         parserContext.m_suggestions.clear();
     }
 };
@@ -351,7 +317,7 @@
 
         parserContext.m_completionIterator = begin;
 
-        const auto& keysNeeded = schema.listKeys(parserContext.currentSchemaPath(), {parserContext.m_curModule, parserContext.m_tmpListName});
+        const auto& keysNeeded = schema.listKeys(dataPathToSchemaPath(parserContext.m_tmpListPath));
         parserContext.m_suggestions = generateMissingKeyCompletionSet(keysNeeded, parserContext.m_tmpListKeys);
     }
 };
@@ -369,7 +335,7 @@
         const auto& dataQuery = parserContext.m_dataquery;
 
         parserContext.m_completionIterator = begin;
-        auto listInstances = dataQuery->listKeys(parserContext.currentDataPath(), {parserContext.m_curModule, parserContext.m_tmpListName});
+        auto listInstances = dataQuery->listKeys(parserContext.m_tmpListPath);
 
         decltype(listInstances) filteredInstances;
         //This filters out instances, which don't correspond to the partial instance we have.
@@ -400,7 +366,7 @@
         const auto& schema = parserContext.m_schema;
 
         parserContext.m_completionIterator = begin;
-        const auto& keysNeeded = schema.listKeys(parserContext.currentSchemaPath(), {parserContext.m_curModule, parserContext.m_tmpListName});
+        const auto& keysNeeded = schema.listKeys(dataPathToSchemaPath(parserContext.m_tmpListPath));
         if (generateMissingKeyCompletionSet(keysNeeded, parserContext.m_tmpListKeys).empty()) {
             parserContext.m_suggestions = {Completion{"]/"}};
         } else {
diff --git a/src/ast_path.cpp b/src/ast_path.cpp
index 1a26a1c..f7c0022 100644
--- a/src/ast_path.cpp
+++ b/src/ast_path.cpp
@@ -61,6 +61,13 @@
 {
 }
 
+dataNode_::dataNode_(boost::optional<module_> module, decltype(m_suffix) node)
+    : m_prefix(module)
+    , m_suffix(node)
+{
+
+}
+
 schemaNode_::schemaNode_(decltype(m_suffix) node)
     : m_suffix(node)
 {
diff --git a/src/ast_path.hpp b/src/ast_path.hpp
index b1b949c..cb84e80 100644
--- a/src/ast_path.hpp
+++ b/src/ast_path.hpp
@@ -97,6 +97,7 @@
 
     dataNode_();
     dataNode_(decltype(m_suffix) node);
+    dataNode_(boost::optional<module_> module, decltype(m_suffix) node);
     dataNode_(module_ module, decltype(m_suffix) node);
     bool operator==(const dataNode_& b) const;
 };
diff --git a/src/data_query.cpp b/src/data_query.cpp
index be979ff..cef7ec3 100644
--- a/src/data_query.cpp
+++ b/src/data_query.cpp
@@ -10,9 +10,9 @@
     m_schema = m_datastore.schema();
 }
 
-std::vector<ListInstance> DataQuery::listKeys(const dataPath_& location, const ModuleNodePair& node) const
+std::vector<ListInstance> DataQuery::listKeys(const dataPath_& listPath) const
 {
-    auto listPath = joinPaths(pathToDataString(location, Prefixes::Always), fullNodeName(location, node));
+    auto listPathString = pathToDataString(listPath, Prefixes::Always);
 
-    return m_datastore.listInstances(listPath);
+    return m_datastore.listInstances(listPathString);
 }
diff --git a/src/data_query.hpp b/src/data_query.hpp
index b2f4a33..7a54a33 100644
--- a/src/data_query.hpp
+++ b/src/data_query.hpp
@@ -24,11 +24,10 @@
 public:
     DataQuery(DatastoreAccess& datastore);
     /*! \brief Lists all possible instances of key/value pairs for a list specified by the arguments.
-     *  \param location Location of the list.
-     *  \param node Name (and optional module name) of the list.
+     *  \param listPath Path to the list (ending with a list_)
      *  \return A vector of maps, which represent the instances. The map is keyed by the name of the list key. The values in the map, are values of the list keys.
      */
-    std::vector<ListInstance> listKeys(const dataPath_& location, const ModuleNodePair& node) const;
+    std::vector<ListInstance> listKeys(const dataPath_& listPath) const;
 
 private:
     DatastoreAccess& m_datastore;
diff --git a/src/parser_context.cpp b/src/parser_context.cpp
index 6b15cb9..723fbae 100644
--- a/src/parser_context.cpp
+++ b/src/parser_context.cpp
@@ -13,14 +13,11 @@
     , m_dataquery(dataQuery)
     , m_curPath(curDir)
 {
-    if (!currentDataPath().m_nodes.empty() && currentDataPath().m_nodes.at(0).m_prefix)
-        m_topLevelModulePresent = true;
 }
 
 void ParserContext::clearPath()
 {
     m_curPath = dataPath_{Scope::Absolute, {}};
-    m_topLevelModulePresent = false;
 }
 
 schemaPath_ ParserContext::currentSchemaPath()
@@ -42,12 +39,9 @@
 
 void ParserContext::pushPathFragment(const dataNode_& node)
 {
-    auto pushNode = [this] (auto& where, const auto& what) {
+    auto pushNode = [] (auto& where, const auto& what) {
         if (std::holds_alternative<nodeup_>(what.m_suffix)) {
             where.m_nodes.pop_back();
-            if (where.m_nodes.empty()) {
-                m_topLevelModulePresent = false;
-            }
         } else {
             where.m_nodes.push_back(what);
         }
@@ -67,13 +61,9 @@
     }
 
     boost::get<schemaPath_>(m_curPath).m_nodes.push_back(node);
-    if (currentSchemaPath().m_nodes.empty()) {
-        m_topLevelModulePresent = false;
-    }
 }
 
 void ParserContext::resetPath()
 {
     m_curPath = m_curPathOrig;
-    m_topLevelModulePresent = !currentDataPath().m_nodes.empty() && currentDataPath().m_nodes.at(0).m_prefix;
 }
diff --git a/src/parser_context.hpp b/src/parser_context.hpp
index 87706a5..88dc71a 100644
--- a/src/parser_context.hpp
+++ b/src/parser_context.hpp
@@ -23,17 +23,17 @@
     const Schema& m_schema;
     const dataPath_ m_curPathOrig;
     const std::shared_ptr<const DataQuery> m_dataquery;
-    boost::optional<std::string> m_curModule;
     std::string m_errorMsg;
-    std::string m_tmpListName;
-    bool m_topLevelModulePresent = false;
 
     struct {
         schemaPath_ m_location;
         ModuleNodePair m_node;
     } m_tmpListKeyLeafPath;
 
+    // When parsing list suffixes, this path is used to store the path of the list whose keys are being parsed.
+    dataPath_ m_tmpListPath;
     std::map<std::string, leaf_data_> m_tmpListKeys;
+
     bool m_errorHandled = false;
     bool m_completing = false;
 
diff --git a/src/path_parser.hpp b/src/path_parser.hpp
index bb622ef..6d9a304 100644
--- a/src/path_parser.hpp
+++ b/src/path_parser.hpp
@@ -166,10 +166,6 @@
 
             auto res = table.parse(begin, end, ctx, rctx, attr);
 
-            if (attr.m_prefix) {
-                parserContext.m_curModule = attr.m_prefix->m_name;
-            }
-
             if (std::holds_alternative<leaf_>(attr.m_suffix)) {
                 parserContext.m_tmpListKeyLeafPath.m_location = parserContext.currentSchemaPath();
                 ModuleNodePair node{attr.m_prefix.flat_map([](const auto& it) {
@@ -181,7 +177,10 @@
 
             if constexpr (std::is_same<attribute_type, dataNode_>()) {
                 if (std::holds_alternative<listElement_>(attr.m_suffix)) {
-                    parserContext.m_tmpListName = std::get<listElement_>(attr.m_suffix).m_name;
+                    parserContext.m_tmpListPath = parserContext.currentDataPath();
+                    auto tmpList = list_{std::get<listElement_>(attr.m_suffix).m_name};
+                    parserContext.m_tmpListPath.m_nodes.push_back(dataNode_{attr.m_prefix, tmpList});
+
                     res = listSuffix.parse(begin, end, ctx, rctx, std::get<listElement_>(attr.m_suffix).m_keys);
 
                     // FIXME: think of a better way to do this, that is, get rid of manual iterator reverting
@@ -219,12 +218,8 @@
 
             if (res) {
                 parserContext.pushPathFragment(attr);
-                parserContext.m_topLevelModulePresent = true;
             }
 
-            if (attr.m_prefix) {
-                parserContext.m_curModule = boost::none;
-            }
             return res;
         }
     }
diff --git a/src/schema.hpp b/src/schema.hpp
index 995d921..42d0728 100644
--- a/src/schema.hpp
+++ b/src/schema.hpp
@@ -63,10 +63,12 @@
     virtual yang::NodeTypes nodeType(const schemaPath_& location, const ModuleNodePair& node) const = 0;
     virtual bool isModule(const std::string& name) const = 0;
     virtual bool listHasKey(const schemaPath_& location, const ModuleNodePair& node, const std::string& key) const = 0;
+    virtual bool listHasKey(const schemaPath_& listPath, const std::string& key) const = 0;
     virtual bool leafIsKey(const std::string& leafPath) const = 0;
     virtual bool isConfig(const std::string& path) const = 0;
     virtual std::optional<std::string> defaultValue(const std::string& leafPath) const = 0;
     virtual const std::set<std::string> listKeys(const schemaPath_& location, const ModuleNodePair& node) const = 0;
+    virtual const std::set<std::string> listKeys(const schemaPath_& listPath) const = 0;
     virtual yang::TypeInfo leafType(const schemaPath_& location, const ModuleNodePair& node) const = 0;
     virtual yang::TypeInfo leafType(const std::string& path) const = 0;
     virtual std::optional<std::string> leafTypeName(const std::string& path) const = 0;
diff --git a/src/static_schema.cpp b/src/static_schema.cpp
index 5fdc42a..203b61e 100644
--- a/src/static_schema.cpp
+++ b/src/static_schema.cpp
@@ -44,6 +44,11 @@
     return list.m_keys.find(key) != list.m_keys.end();
 }
 
+bool StaticSchema::listHasKey(const schemaPath_& listPath, const std::string& key) const
+{
+    return listKeys(listPath).count(key);
+}
+
 const std::set<std::string> StaticSchema::listKeys(const schemaPath_& location, const ModuleNodePair& node) const
 {
     std::string locationString = pathToSchemaString(location, Prefixes::Always);
@@ -54,6 +59,23 @@
     return list.m_keys;
 }
 
+std::string lastNodeOfSchemaPath(const std::string& path)
+{
+    std::string res = path;
+    if (auto pos = res.find_last_of('/'); pos != res.npos) {
+        res.erase(0, pos + 1);
+    }
+    return res;
+}
+
+const std::set<std::string> StaticSchema::listKeys(const schemaPath_& listPath) const
+{
+    auto listPathString = pathToSchemaString(listPath, Prefixes::Always);
+    const auto& child = children(stripLastNodeFromPath(listPathString)).at(lastNodeOfSchemaPath(listPathString));
+    const auto& list = std::get<yang::list>(child.m_nodeType);
+    return list.m_keys;
+}
+
 void StaticSchema::addList(const std::string& location, const std::string& name, const std::set<std::string>& keys)
 {
     m_nodes.at(location).emplace(name, NodeInfo{yang::list{keys}, yang::AccessType::Writable});
@@ -106,15 +128,6 @@
     }
 }
 
-std::string lastNodeOfSchemaPath(const std::string& path)
-{
-    std::string res = path;
-    if (auto pos = res.find_last_of('/'); pos != res.npos) {
-        res.erase(0, pos + 1);
-    }
-    return res;
-}
-
 yang::TypeInfo StaticSchema::leafType(const schemaPath_& location, const ModuleNodePair& node) const
 {
     std::string locationString = pathToSchemaString(location, Prefixes::Always);
diff --git a/src/static_schema.hpp b/src/static_schema.hpp
index 0426f02..09b6271 100644
--- a/src/static_schema.hpp
+++ b/src/static_schema.hpp
@@ -63,10 +63,12 @@
     yang::NodeTypes nodeType(const schemaPath_& location, const ModuleNodePair& node) const override;
     bool isModule(const std::string& name) const override;
     bool listHasKey(const schemaPath_& location, const ModuleNodePair& node, const std::string& key) const override;
+    bool listHasKey(const schemaPath_& listPath, const std::string& key) const override;
     bool leafIsKey(const std::string& leafPath) const override;
     bool isConfig(const std::string& leafPath) const override;
     std::optional<std::string> defaultValue(const std::string& leafPath) const override;
     const std::set<std::string> listKeys(const schemaPath_& location, const ModuleNodePair& node) const override;
+    const std::set<std::string> listKeys(const schemaPath_& listPath) const override;
     yang::TypeInfo leafType(const schemaPath_& location, const ModuleNodePair& node) const override;
     yang::TypeInfo leafType(const std::string& path) const override;
     std::optional<std::string> leafTypeName(const std::string& path) const override;
diff --git a/src/yang_schema.cpp b/src/yang_schema.cpp
index 48b3724..b44ef5f 100644
--- a/src/yang_schema.cpp
+++ b/src/yang_schema.cpp
@@ -101,6 +101,12 @@
     return keys.find(key) != keys.end();
 }
 
+bool YangSchema::listHasKey(const schemaPath_& listPath, const std::string& key) const
+{
+    const auto keys = listKeys(listPath);
+    return keys.find(key) != keys.end();
+}
+
 bool YangSchema::leafIsKey(const std::string& leafPath) const
 {
     auto node = getSchemaNode(leafPath);
@@ -141,17 +147,36 @@
     return impl_getSchemaNode(absPath);
 }
 
-const std::set<std::string> YangSchema::listKeys(const schemaPath_& location, const ModuleNodePair& node) const
+libyang::S_Schema_Node YangSchema::getSchemaNode(const schemaPath_& listPath) const
+{
+    std::string absPath = pathToSchemaString(listPath, Prefixes::Always);
+    return impl_getSchemaNode(absPath);
+}
+
+namespace {
+const std::set<std::string> impl_listKeys(const libyang::S_Schema_Node_List& list)
 {
     std::set<std::string> keys;
-    if (!isList(location, node))
-        return keys;
-    libyang::Schema_Node_List list(getSchemaNode(location, node));
-    const auto& keysVec = list.keys();
+    const auto& keysVec = list->keys();
 
     std::transform(keysVec.begin(), keysVec.end(), std::inserter(keys, keys.begin()), [](const auto& it) { return it->name(); });
     return keys;
 }
+}
+
+const std::set<std::string> YangSchema::listKeys(const schemaPath_& location, const ModuleNodePair& node) const
+{
+    if (!isList(location, node))
+        return {};
+    auto list = std::make_shared<libyang::Schema_Node_List>(getSchemaNode(location, node));
+    return impl_listKeys(list);
+}
+
+const std::set<std::string> YangSchema::listKeys(const schemaPath_& listPath) const
+{
+    auto list = std::make_shared<libyang::Schema_Node_List>(getSchemaNode(listPath));
+    return impl_listKeys(list);
+}
 
 namespace {
 std::set<enum_> enumValues(const libyang::S_Type& typeArg)
diff --git a/src/yang_schema.hpp b/src/yang_schema.hpp
index 5b2db04..aed1a5d 100644
--- a/src/yang_schema.hpp
+++ b/src/yang_schema.hpp
@@ -35,10 +35,12 @@
     yang::NodeTypes nodeType(const schemaPath_& location, const ModuleNodePair& node) const override;
     bool isModule(const std::string& name) const override;
     bool listHasKey(const schemaPath_& location, const ModuleNodePair& node, const std::string& key) const override;
+    bool listHasKey(const schemaPath_& listPath, const std::string& key) const override;
     bool leafIsKey(const std::string& leafPath) const override;
     bool isConfig(const std::string& path) const override;
     std::optional<std::string> defaultValue(const std::string& leafPath) const override;
     const std::set<std::string> listKeys(const schemaPath_& location, const ModuleNodePair& node) const override;
+    const std::set<std::string> listKeys(const schemaPath_& listPath) const override;
     yang::TypeInfo leafType(const schemaPath_& location, const ModuleNodePair& node) const override;
     yang::TypeInfo leafType(const std::string& path) const override;
     /** @brief If the leaf type is a typedef, returns the typedef name. */
@@ -74,10 +76,11 @@
     yang::TypeInfo impl_leafType(const std::shared_ptr<libyang::Schema_Node>& node) const;
     std::set<std::string> modules() const;
 
-    /** @short Returns a set of nodes, that match the location and name criteria. */
 
     /** @short Returns a single Schema_Node if the criteria matches only one, otherwise nullptr. */
     std::shared_ptr<libyang::Schema_Node> getSchemaNode(const std::string& node) const;
+    /** @short Returns a single Schema_Node if the criteria matches only one, otherwise nullptr. */
+    std::shared_ptr<libyang::Schema_Node> getSchemaNode(const schemaPath_& listPath) const;
 
     /** @short Returns a single Schema_Node if the criteria matches only one, otherwise nullptr. */
     std::shared_ptr<libyang::Schema_Node> getSchemaNode(const schemaPath_& location, const ModuleNodePair& node) const;
diff --git a/tests/data_query.cpp b/tests/data_query.cpp
index 28c51e8..410c7c4 100644
--- a/tests/data_query.cpp
+++ b/tests/data_query.cpp
@@ -56,9 +56,8 @@
 
     SECTION("listKeys")
     {
-        dataPath_ location;
-        location.m_scope = Scope::Absolute;
-        ModuleNodePair node;
+        dataPath_ listPath;
+        listPath.m_scope = Scope::Absolute;
         std::vector<std::map<std::string, leaf_data_>> expected;
 
         SECTION("example-schema:person")
@@ -66,8 +65,7 @@
             datastore.createListInstance("/example-schema:person[name='Vaclav']");
             datastore.createListInstance("/example-schema:person[name='Tomas']");
             datastore.createListInstance("/example-schema:person[name='Jan Novak']");
-            node.first = "example-schema";
-            node.second = "person";
+            listPath.m_nodes.push_back(dataNode_{{"example-schema"}, list_{"person"}});
             expected = {
                 {{"name", std::string{"Jan Novak"}}},
                 {{"name", std::string{"Tomas"}}},
@@ -77,8 +75,7 @@
 
         SECTION("example-schema:person - no instances")
         {
-            node.first = "example-schema";
-            node.second = "person";
+            listPath.m_nodes.push_back(dataNode_{{"example-schema"}, list_{"person"}});
             expected = {
             };
         }
@@ -88,8 +85,7 @@
             datastore.createListInstance("/example-schema:selectedNumbers[value='45']");
             datastore.createListInstance("/example-schema:selectedNumbers[value='99']");
             datastore.createListInstance("/example-schema:selectedNumbers[value='127']");
-            node.first = "example-schema";
-            node.second = "selectedNumbers";
+            listPath.m_nodes.push_back(dataNode_{{"example-schema"}, list_{"selectedNumbers"}});
             expected = {
                 {{"value", int8_t{127}}},
                 {{"value", int8_t{45}}},
@@ -102,8 +98,7 @@
             datastore.createListInstance("/example-schema:animalWithColor[name='Dog'][color='brown']");
             datastore.createListInstance("/example-schema:animalWithColor[name='Dog'][color='white']");
             datastore.createListInstance("/example-schema:animalWithColor[name='Cat'][color='grey']");
-            node.first = "example-schema";
-            node.second = "animalWithColor";
+            listPath.m_nodes.push_back(dataNode_{{"example-schema"}, list_{"animalWithColor"}});
             expected = {
                 {{"name", std::string{"Cat"}}, {"color", std::string{"grey"}}},
                 {{"name", std::string{"Dog"}}, {"color", std::string{"brown"}}},
@@ -114,8 +109,7 @@
         SECTION("example-schema:animalWithColor - quotes in values")
         {
             datastore.createListInstance("/example-schema:animalWithColor[name='D\"o\"g'][color=\"b'r'own\"]");
-            node.first = "example-schema";
-            node.second = "animalWithColor";
+            listPath.m_nodes.push_back(dataNode_{{"example-schema"}, list_{"animalWithColor"}});
             expected = {
                 {{"name", std::string{"D\"o\"g"}}, {"color", std::string{"b'r'own"}}}
             };
@@ -126,8 +120,7 @@
             datastore.createListInstance("/example-schema:ports[name='A']");
             datastore.createListInstance("/example-schema:ports[name='B']");
             datastore.createListInstance("/example-schema:ports[name='E']");
-            node.first = "example-schema";
-            node.second = "ports";
+            listPath.m_nodes.push_back(dataNode_{{"example-schema"}, list_{"ports"}});
             expected = {
                 {{"name", enum_{"A"}}},
                 {{"name", enum_{"B"}}},
@@ -151,8 +144,7 @@
 
             SECTION("outer list")
             {
-                node.first = "example-schema";
-                node.second = "org";
+                listPath.m_nodes.push_back(dataNode_{{"example-schema"}, list_{"org"}});
                 expected = {
                     {{"department", std::string{"accounting"}}},
                     {{"department", std::string{"sales"}}},
@@ -164,7 +156,6 @@
             {
                 listElement_ list;
                 list.m_name = "org";
-                node.second = "people";
                 SECTION("accounting department")
                 {
                     list.m_keys = {
@@ -193,7 +184,8 @@
                     expected = {
                     };
                 }
-                location.m_nodes.push_back(dataNode_{{"example-schema"}, list});
+                listPath.m_nodes.push_back(dataNode_{{"example-schema"}, list});
+                listPath.m_nodes.push_back(dataNode_{list_{"people"}});
             }
 
             SECTION("THREE MF NESTED LISTS")
@@ -205,7 +197,6 @@
                 };
 
                 listElement_ listPeople;
-                node.second = "computers";
 
                 SECTION("alice computers")
                 {
@@ -232,9 +223,9 @@
                     };
                 }
 
-                location.m_nodes.push_back(dataNode_{{"example-schema"}, listOrg});
-                location.m_nodes.push_back(dataNode_{listPeople});
-
+                listPath.m_nodes.push_back(dataNode_{{"example-schema"}, listOrg});
+                listPath.m_nodes.push_back(dataNode_{listPeople});
+                listPath.m_nodes.push_back(dataNode_{list_{"computers"}});
             }
         }
 
@@ -243,9 +234,8 @@
             datastore.createListInstance("/other-module:parking-lot/example-schema:cars[id='1']");
             datastore.createListInstance("/other-module:parking-lot/example-schema:cars[id='2']");
 
-            location.m_nodes.push_back(dataNode_{{"other-module"}, container_{"parking-lot"}});
-            node.first = "example-schema";
-            node.second = "cars";
+            listPath.m_nodes.push_back(dataNode_{{"other-module"}, container_{"parking-lot"}});
+            listPath.m_nodes.push_back(dataNode_{{"example-schema"}, list_{"cars"}});
             expected = {
                 {{"id", int32_t{1}}},
                 {{"id", int32_t{2}}},
@@ -255,7 +245,7 @@
 
         datastore.commitChanges();
         std::sort(expected.begin(), expected.end());
-        auto keys = dataquery.listKeys(location, node);
+        auto keys = dataquery.listKeys(listPath);
         std::sort(keys.begin(), keys.end());
         REQUIRE(keys == expected);
     }
diff --git a/tests/list_manipulation.cpp b/tests/list_manipulation.cpp
index 974d806..d2ad5f9 100644
--- a/tests/list_manipulation.cpp
+++ b/tests/list_manipulation.cpp
@@ -7,6 +7,7 @@
 
 #include "trompeloeil_doctest.hpp"
 #include "parser.hpp"
+#include "pretty_printers.hpp"
 #include "static_schema.hpp"
 
 TEST_CASE("list manipulation")
@@ -14,10 +15,17 @@
     using namespace std::string_literals;
     auto schema = std::make_shared<StaticSchema>();
     schema->addModule("mod");
+    schema->addModule("other");
     schema->addList("/", "mod:list", {"number"});
     schema->addLeaf("/mod:list", "mod:number", yang::Int32{});
     schema->addLeaf("/mod:list", "mod:leafInList", yang::String{});
     schema->addLeafList("/", "mod:addresses", yang::String{});
+    schema->addIdentity(std::nullopt, identityRef_{"other", "deptypes"});
+    schema->addIdentity(identityRef_{"other", "deptypes"}, identityRef_{"other", "engineering"});
+    schema->addList("/", "mod:company", {"department"});
+    schema->addLeaf("/mod:company", "mod:department", schema->validIdentities("other", "deptypes"));
+    schema->addList("/mod:company", "mod:inventory", {"id"});
+    schema->addLeaf("/mod:company/mod:inventory", "mod:id", yang::Int32{});
     Parser parser(schema);
     std::string input;
     std::ostringstream errorStream;
@@ -33,6 +41,17 @@
             expectedPath.m_nodes.push_back(dataNode_{module_{"mod"}, listElement_("list", keys)});
         }
 
+        SECTION("mod:company[department=other:engineering]/inventory[id=1337]")
+        {
+            input = "mod:company[department=other:engineering]/inventory[id=1337]";
+            auto keys = std::map<std::string, leaf_data_>{
+                {"department", identityRef_{"other", "engineering"}}};
+            expectedPath.m_nodes.push_back(dataNode_{module_{"mod"}, listElement_("company", keys)});
+            keys = std::map<std::string, leaf_data_>{
+                {"id", int32_t{1337}}};
+            expectedPath.m_nodes.push_back(dataNode_{listElement_("inventory", keys)});
+        }
+
         SECTION("create mod:addresses['0.0.0.0']")
         {
             input = "mod:addresses['0.0.0.0']";
diff --git a/tests/ls_interpreter.cpp b/tests/ls_interpreter.cpp
index a78194e..4a60762 100644
--- a/tests/ls_interpreter.cpp
+++ b/tests/ls_interpreter.cpp
@@ -26,9 +26,11 @@
     IMPLEMENT_CONST_MOCK1(leafTypeName);
     IMPLEMENT_CONST_MOCK1(isModule);
     IMPLEMENT_CONST_MOCK1(leafrefPath);
-    IMPLEMENT_CONST_MOCK3(listHasKey);
+    MAKE_CONST_MOCK3(listHasKey, bool(const schemaPath_& location, const ModuleNodePair& node, const std::string& key), override);
+    MAKE_CONST_MOCK2(listHasKey, bool(const schemaPath_& listPath, const std::string& key), override);
     IMPLEMENT_CONST_MOCK1(leafIsKey);
-    IMPLEMENT_CONST_MOCK2(listKeys);
+    MAKE_CONST_MOCK2(listKeys, const std::set<std::string>(const schemaPath_& location, const ModuleNodePair& node), override);
+    MAKE_CONST_MOCK1(listKeys, const std::set<std::string>(const schemaPath_& listPath), override);
     MAKE_CONST_MOCK1(nodeType, yang::NodeTypes(const std::string&), override);
     MAKE_CONST_MOCK2(nodeType, yang::NodeTypes(const schemaPath_&, const ModuleNodePair&), override);
     IMPLEMENT_CONST_MOCK1(status);
diff --git a/tests/pretty_printers.hpp b/tests/pretty_printers.hpp
index 7f59adc..429a0e9 100644
--- a/tests/pretty_printers.hpp
+++ b/tests/pretty_printers.hpp
@@ -122,6 +122,11 @@
     return s;
 }
 
+std::ostream& operator<<(std::ostream& s, const create_& create)
+{
+    s << "\nls_ {\n    " << create.m_path << "}\n";
+    return s;
+}
 
 std::ostream& operator<<(std::ostream& s, const ls_& ls)
 {