diff --git a/tests/leaf_editing.cpp b/tests/leaf_editing.cpp
index ad88a99..64abbf2 100644
--- a/tests/leaf_editing.cpp
+++ b/tests/leaf_editing.cpp
@@ -50,26 +50,26 @@
     schema->addList("/", "mod:list", {"number"});
     schema->addLeaf("/mod:list", "mod:number", yang::Int32{});
     schema->addLeaf("/mod:list", "mod:leafInList", yang::String{});
-    schema->addLeaf("/", "mod:refToString", yang::LeafRef{"/mod:leafString", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:leafString"))});
-    schema->addLeaf("/", "mod:refToInt8", yang::LeafRef{"/mod:leafInt8", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:leafInt8"))});
-    schema->addLeaf("/", "mod:refToLeafInCont", yang::LeafRef{"/mod:contA/identInCont", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:contA/mod:identInCont"))});
-    schema->addLeaf("/", "mod:intOrString", yang::Union{{yang::Int32{}, yang::String{}}});
-    schema->addLeaf("/", "mod:twoInts", yang::Union{{yang::Uint8{}, yang::Int16{}}});
+    schema->addLeaf("/", "mod:refToString", yang::LeafRef{"/mod:leafString", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:leafString"))});
+    schema->addLeaf("/", "mod:refToInt8", yang::LeafRef{"/mod:leafInt8", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:leafInt8"))});
+    schema->addLeaf("/", "mod:refToLeafInCont", yang::LeafRef{"/mod:contA/identInCont", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:contA/mod:identInCont"))});
+    schema->addLeaf("/", "mod:intOrString", yang::Union{{yang::TypeInfo{yang::Int32{}}, yang::TypeInfo{yang::String{}}}});
+    schema->addLeaf("/", "mod:twoInts", yang::Union{{yang::TypeInfo{yang::Uint8{}}, yang::TypeInfo{yang::Int16{}}}});
     schema->addLeaf("/", "mod:unionStringEnumLeafref", yang::Union{{
-        yang::String{},
-        createEnum({"foo", "bar"}),
-        yang::LeafRef{"/mod:leafEnum", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:leafEnum"))}
+            yang::LeafDataType{yang::String{}},
+        yang::LeafDataType{createEnum({"foo", "bar"})},
+        yang::LeafDataType{yang::LeafRef{"/mod:leafEnum", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:leafEnum"))}}
     }});
 
     schema->addList("/", "mod:portSettings", {"port"});
     schema->addLeaf("/mod:portSettings", "mod:port", createEnum({"eth0", "eth1", "eth2"}));
     schema->addList("/", "mod:portMapping", {"port"});
     schema->addLeaf("/mod:portMapping", "mod:port", createEnum({"utf1", "utf2", "utf3"}));
-    schema->addLeaf("/", "mod:activeMappedPort", yang::LeafRef{"/mod:portMapping/mod:port", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:portMapping/mod:port"))});
+    schema->addLeaf("/", "mod:activeMappedPort", yang::LeafRef{"/mod:portMapping/mod:port", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:portMapping/mod:port"))});
     schema->addLeaf("/", "mod:activePort", yang::Union{{
-        createEnum({"wlan0", "wlan1"}),
-        yang::LeafRef{"/mod:portSettings/mod:port", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:portSettings/mod:port"))},
-        yang::LeafRef{"/mod:activeMappedPort", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:activeMappedPort"))},
+        yang::TypeInfo{createEnum({"wlan0", "wlan1"})},
+        yang::TypeInfo{yang::LeafRef{"/mod:portSettings/mod:port", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:portSettings/mod:port"))}},
+        yang::TypeInfo{yang::LeafRef{"/mod:activeMappedPort", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:activeMappedPort"))}},
     }});
 
     Parser parser(schema);
diff --git a/tests/pretty_printers.hpp b/tests/pretty_printers.hpp
index e5b6659..98638ea 100644
--- a/tests/pretty_printers.hpp
+++ b/tests/pretty_printers.hpp
@@ -64,14 +64,22 @@
         s << "}";
     }
     if (std::holds_alternative<yang::LeafRef>(type)) {
-        s << "{" << std::get<yang::LeafRef>(type).m_targetXPath << "," << *std::get<yang::LeafRef>(type).m_targetType  << "}";
+        s << "{" << std::get<yang::LeafRef>(type).m_targetXPath << "," << std::get<yang::LeafRef>(type).m_targetType->m_type  << "}";
     }
     if (std::holds_alternative<yang::Union>(type)) {
         s << "{" << std::endl;
         auto types = std::get<yang::Union>(type).m_unionTypes;
-        std::copy(types.begin(), types.end(), std::experimental::make_ostream_joiner(s, ",\n"));
+        std::transform(types.begin(), types.end(), std::experimental::make_ostream_joiner(s, ",\n"), [] (const auto& type) {
+            return type.m_type;
+        });
     }
     s << std::endl;
     return s;
 }
+
+std::ostream& operator<<(std::ostream& s, const yang::TypeInfo& type)
+{
+    s << type.m_type << (type.m_units ? " units: " + *type.m_units : "");
+    return s;
+}
 }
diff --git a/tests/utils.cpp b/tests/utils.cpp
index cea698b..509c2fe 100644
--- a/tests/utils.cpp
+++ b/tests/utils.cpp
@@ -89,10 +89,10 @@
         SECTION("union")
         {
             type = yang::Union{{
-                yang::String{},
-                createEnum({"foo", "bar"}),
-                yang::Int8{},
-                yang::Int64{},
+                yang::TypeInfo{yang::String{}},
+                yang::TypeInfo{createEnum({"foo", "bar"})},
+                yang::TypeInfo{yang::Int8{}},
+                yang::TypeInfo{yang::Int64{}},
             }};
             expected = "a string, an enum, an 8-bit integer, a 64-bit integer";
         }
diff --git a/tests/yang.cpp b/tests/yang.cpp
index a9fb65e..941586d 100644
--- a/tests/yang.cpp
+++ b/tests/yang.cpp
@@ -695,7 +695,7 @@
                 node.second = "activeNumber";
                 type.emplace<yang::LeafRef>(
                     "/example-schema:_list/number",
-                    std::make_unique<yang::LeafDataType>(ys.leafType("/example-schema:_list/number"))
+                    std::make_unique<yang::TypeInfo>(ys.leafType("/example-schema:_list/number"))
                 );
             }
 
@@ -718,17 +718,18 @@
                 }();
 
                 type = yang::Union{{
-                    createEnum({"wlan0", "wlan1"}),
-                    yang::LeafRef{
+                    yang::TypeInfo{createEnum({"wlan0", "wlan1"})},
+                    yang::TypeInfo{yang::LeafRef{
                         "/example-schema:portSettings/port",
-                        std::make_unique<yang::LeafDataType>(createEnum({"eth0", "eth1", "eth2"}))
-                    },
-                    yang::LeafRef{
+                        std::make_unique<yang::TypeInfo>(createEnum({"eth0", "eth1", "eth2"}))
+                    }},
+                    yang::TypeInfo{yang::LeafRef{
                         "/example-schema:activeMappedPort",
-                        std::make_unique<yang::LeafDataType>(yang::LeafRef{
+                        std::make_unique<yang::TypeInfo>(yang::LeafRef{
                                 "/example-schema:portMapping/port",
-                                std::make_unique<yang::LeafDataType>(enums)
-                        })},
+                                std::make_unique<yang::TypeInfo>(enums)
+                        })
+                    }},
                 }};
             }
 
@@ -841,37 +842,42 @@
 
         SECTION("units")
         {
-            std::optional<std::string> expected;
+            yang::LeafDataType expectedType;
+            std::optional<std::string> expectedUnits;
             SECTION("length")
             {
                 path.m_nodes.push_back(schemaNode_(module_{"example-schema"}, leaf_("length")));
-                expected = "m";
+                expectedType.emplace<yang::Int32>();
+                expectedUnits = "m";
             }
 
             SECTION("wavelength")
             {
                 path.m_nodes.push_back(schemaNode_(module_{"example-schema"}, leaf_("wavelength")));
-                expected = "nm";
+                expectedType.emplace<yang::Decimal>();
+                expectedUnits = "nm";
             }
 
             SECTION("leafInt32")
             {
                 path.m_nodes.push_back(schemaNode_(module_{"example-schema"}, leaf_("leafInt32")));
+                expectedType.emplace<yang::Int32>();
             }
 
             SECTION("duration")
             {
                 path.m_nodes.push_back(schemaNode_(module_{"example-schema"}, leaf_("duration")));
-                expected = "s";
+                expectedType.emplace<yang::Int32>();
+                expectedUnits = "s";
             }
 
             SECTION("another-duration")
             {
                 path.m_nodes.push_back(schemaNode_(module_{"example-schema"}, leaf_("another-duration")));
-                expected = "vt";
+                expectedType.emplace<yang::Int32>();
+                expectedUnits = "vt";
             }
-
-            REQUIRE(ys.units(pathToSchemaString(path, Prefixes::WhenNeeded)) == expected);
+            REQUIRE(ys.leafType(pathToSchemaString(path, Prefixes::WhenNeeded)) == yang::TypeInfo{expectedType, expectedUnits});
         }
 
         SECTION("nodeType")