diff --git a/src/class_diagram/visitor/translation_unit_visitor.cc b/src/class_diagram/visitor/translation_unit_visitor.cc
index 92595dc7..7c911349 100644
--- a/src/class_diagram/visitor/translation_unit_visitor.cc
+++ b/src/class_diagram/visitor/translation_unit_visitor.cc
@@ -1162,9 +1162,6 @@ void translation_unit_visitor::
 {
     auto template_params = cx::util::parse_unexposed_template_params(
         type_name, [this](const std::string &t) {
-            //            auto full_type = ctx.get_name_with_namespace(t);
-            //            if (full_type.has_value())
-            //                return full_type.value().to_string();
             return t;
         });
 
@@ -1177,27 +1174,6 @@ void translation_unit_visitor::
     for (auto &r : relationships) {
         c.add_relationship({std::get<1>(r), std::get<0>(r)});
     }
-
-    //    const auto &primary_template_ref =
-    //        static_cast<const cppast::cpp_class_template &>(
-    //            tspec.value().primary_template().get(ctx.entity_index())[0].get())
-    //            .class_();
-
-    //    if (primary_template_ref.user_data()) {
-    //        auto base_template_full_name =
-    //            static_cast<const char *>(primary_template_ref.user_data());
-    //        LOG_DBG("Primary template ref set to: {}",
-    //        base_template_full_name);
-    //        // Add template specialization/instantiation
-    //        // relationship
-    //        c.add_relationship(
-    //            {relationship_t::kInstantiation, base_template_full_name});
-    //    }
-    //    else {
-    //        LOG_DBG(
-    //            "No user data for base template {}",
-    //            primary_template_ref.name());
-    //    }
 }
 
 bool translation_unit_visitor::find_relationships_in_unexposed_template_params(
diff --git a/src/common/clang_utils.cc b/src/common/clang_utils.cc
index 208909a3..74bfaa45 100644
--- a/src/common/clang_utils.cc
+++ b/src/common/clang_utils.cc
@@ -31,4 +31,10 @@ template <> id_t to_id(const clang::TemplateSpecializationType &t)
 {
     return t.getTemplateName().getAsTemplateDecl()->getID();
 }
+
+template <> id_t to_id(const std::filesystem::path &file)
+{
+    return std::hash<std::string>{}(file.lexically_normal()) >> 3;
+}
+
 }
diff --git a/src/common/clang_utils.h b/src/common/clang_utils.h
index 1adcdf8a..9d842c76 100644
--- a/src/common/clang_utils.h
+++ b/src/common/clang_utils.h
@@ -23,6 +23,7 @@
 #include <clang/AST/RecursiveASTVisitor.h>
 
 #include <string>
+#include <filesystem>
 
 namespace clang {
 class NamespaceDecl;
@@ -49,4 +50,5 @@ template <> id_t to_id(const clang::EnumType &type);
 
 template <> id_t to_id(const clang::TemplateSpecializationType &type);
 
+template <> id_t to_id(const std::filesystem::path &type);
 }
diff --git a/src/common/generators/plantuml/generator.h b/src/common/generators/plantuml/generator.h
index be87724b..2a699c17 100644
--- a/src/common/generators/plantuml/generator.h
+++ b/src/common/generators/plantuml/generator.h
@@ -22,7 +22,6 @@
 #include "util/error.h"
 #include "util/util.h"
 
-//#include <cppast/libclang_parser.hpp>
 #include <clang/Frontend/CompilerInstance.h>
 #include <clang/Tooling/CompilationDatabase.h>
 #include <clang/Tooling/Tooling.h>
@@ -262,8 +261,6 @@ public:
 
     virtual void HandleTranslationUnit(clang::ASTContext &ast_context)
     {
-        //        const auto* tud = ast_context.getTranslationUnitDecl();
-        ////        tud->dump();
         visitor_.TraverseDecl(ast_context.getTranslationUnitDecl());
         visitor_.finalize();
     }
@@ -280,14 +277,31 @@ public:
     {
     }
 
-    virtual std::unique_ptr<clang::ASTConsumer> CreateASTConsumer(
-        clang::CompilerInstance &CI, clang::StringRef file)
+    std::unique_ptr<clang::ASTConsumer> CreateASTConsumer(
+        clang::CompilerInstance &CI, clang::StringRef file) override
     {
         return std::make_unique<
             diagram_ast_consumer<DiagramModel, DiagramConfig, DiagramVisitor>>(
             CI, diagram_, config_);
     }
 
+protected:
+    bool BeginSourceFileAction(clang::CompilerInstance &ci) override
+    {
+        if constexpr (std::is_same_v<DiagramModel,
+                          clanguml::include_diagram::model::diagram>) {
+            auto find_includes_callback =
+                std::make_unique<typename DiagramVisitor::include_visitor>(
+                    ci.getSourceManager(), diagram_, config_);
+
+            clang::Preprocessor &pp = ci.getPreprocessor();
+
+            pp.addPPCallbacks(std::move(find_includes_callback));
+        }
+
+        return true;
+    }
+
 private:
     DiagramModel &diagram_;
     const DiagramConfig &config_;
@@ -338,28 +352,12 @@ std::unique_ptr<DiagramModel> generate(
             std::back_inserter(translation_units));
     }
 
-    //    DiagramVisitor visitor(db, *diagram, config);
-
     clang::tooling::ClangTool clang_tool(db, translation_units);
     auto action_factory =
         std::make_unique<diagram_action_visitor_factory<DiagramModel,
             DiagramConfig, DiagramVisitor>>(*diagram, config);
     clang_tool.run(action_factory.get());
 
-    /*
-    cppast::cpp_entity_index idx;
-    auto logger =
-        verbose ? cppast::default_logger() : cppast::default_quiet_logger();
-    cppast::simple_file_parser<cppast::libclang_parser> parser{
-        type_safe::ref(idx), std::move(logger)};
-
-    // Process all matching translation units
-    DiagramVisitor ctx(idx, *diagram, config);
-    cppast::parse_files(parser, translation_units, db);
-    for (auto &file : parser.files())
-        ctx(file);
-    */
-
     diagram->set_complete(true);
 
     return diagram;
diff --git a/src/common/model/diagram_filter.cc b/src/common/model/diagram_filter.cc
index 1fe1f568..17473065 100644
--- a/src/common/model/diagram_filter.cc
+++ b/src/common/model/diagram_filter.cc
@@ -357,6 +357,10 @@ tvl::value_t paths_filter::match(
         return {};
     }
 
+    // Matching source paths doesn't make sens if they are not absolute
+    if(!p.is_absolute())
+        return {};
+
     auto pp = p.fs_path(root_);
     for (const auto &path : paths_) {
         if (util::starts_with(pp, path))
diff --git a/src/common/model/diagram_filter.h b/src/common/model/diagram_filter.h
index 4e084d6e..aa8ddd9c 100644
--- a/src/common/model/diagram_filter.h
+++ b/src/common/model/diagram_filter.h
@@ -238,8 +238,9 @@ private:
         // of matching elements
         for (const auto &template_root : roots_) {
             auto template_ref = detail::get<ElementT>(cd, template_root);
-            if (template_ref.has_value())
+            if (template_ref.has_value()) {
                 matching_elements_.emplace(template_ref.value());
+            }
         }
 
         assert(roots_.empty() == matching_elements_.empty());
diff --git a/src/common/model/source_file.h b/src/common/model/source_file.h
index a99599af..3c0d5c41 100644
--- a/src/common/model/source_file.h
+++ b/src/common/model/source_file.h
@@ -17,6 +17,7 @@
  */
 #pragma once
 
+#include "common/clang_utils.h"
 #include "common/model/diagram_element.h"
 #include "common/model/nested_trait.h"
 #include "common/model/path.h"
@@ -51,17 +52,20 @@ class source_file
 public:
     source_file() = default;
 
-    source_file(const std::filesystem::path &p)
+    explicit source_file(const std::filesystem::path &p)
     {
         set_path({p.parent_path().string()});
         set_name(p.filename());
         is_absolute_ = p.is_absolute();
+        set_id(common::to_id(p));
     }
 
     void set_path(const filesystem_path &p) { path_ = p; }
 
     void set_absolute() { is_absolute_ = true; }
 
+    bool is_absolute() const { return is_absolute_; }
+
     void set_type(source_file_t type) { type_ = type; }
 
     source_file_t type() const { return type_; }
@@ -71,6 +75,12 @@ public:
     source_file &operator=(const source_file &) = delete;
     source_file &operator=(source_file &&) = delete;
 
+    bool operator==(const source_file &right) const
+    {
+        return (path_ == right.path_) && (name() == right.name()) &&
+            (type_ == right.type_);
+    }
+
     const filesystem_path &path() const { return path_; }
 
     std::string full_name(bool /*relative*/) const override
diff --git a/src/include_diagram/generators/plantuml/include_diagram_generator.cc b/src/include_diagram/generators/plantuml/include_diagram_generator.cc
index 1543ff6a..5c7e2695 100644
--- a/src/include_diagram/generators/plantuml/include_diagram_generator.cc
+++ b/src/include_diagram/generators/plantuml/include_diagram_generator.cc
@@ -37,33 +37,36 @@ void generator::generate_relationships(
 
     namespace plantuml_common = clanguml::common::generators::plantuml;
 
-    //    if (f.type() == common::model::source_file_t::kDirectory) {
-    //        util::for_each(f, [this, &ostr](const auto &file) {
-    //            generate_relationships(
-    //                dynamic_cast<const source_file &>(*file), ostr);
-    //        });
-    //    }
-    //    else {
-    //        util::for_each_if(
-    //            f.relationships(),
-    //            [this](const auto &r) {
-    //                return m_model.should_include(r.type()) &&
-    //                    util::contains(m_generated_aliases, r.destination());
-    //            },
-    //            [&f, &ostr](const auto &r) {
-    //                ostr << f.alias() << " "
-    //                     << plantuml_common::to_plantuml(r.type(), r.style())
-    //                     << " "
-    //                     << r.destination() << '\n';
-    //            });
-    //    }
+    if (f.type() == common::model::source_file_t::kDirectory) {
+        util::for_each(f, [this, &ostr](const auto &file) {
+            generate_relationships(
+                dynamic_cast<const source_file &>(*file), ostr);
+        });
+    }
+    else {
+        util::for_each_if(
+            f.relationships(),
+            [this](const auto &r) {
+                return m_model.should_include(r.type()) &&
+                    util::contains(m_generated_aliases,
+                        m_model.get(r.destination()).value().get().alias());
+            },
+            [&f, &ostr, this](const auto &r) {
+                ostr << f.alias() << " "
+                     << plantuml_common::to_plantuml(r.type(), r.style()) << " "
+                     << m_model.get(r.destination()).value().get().alias()
+                     << '\n';
+            });
+    }
 }
 
 void generator::generate(const source_file &f, std::ostream &ostr) const
 {
-    LOG_DBG("Generating source_file {}", f.name());
+
 
     if (f.type() == common::model::source_file_t::kDirectory) {
+        LOG_DBG("Generating directory {}", f.name());
+
         ostr << "folder \"" << f.name();
         ostr << "\" as " << f.alias();
         ostr << " {\n";
@@ -77,6 +80,8 @@ void generator::generate(const source_file &f, std::ostream &ostr) const
         m_generated_aliases.emplace(f.alias());
     }
     else {
+        LOG_DBG("Generating file {}", f.name());
+
         if (m_model.should_include(f)) {
             ostr << "file \"" << f.name() << "\" as " << f.alias();
 
@@ -95,14 +100,14 @@ void generator::generate(std::ostream &ostr) const
 {
     ostr << "@startuml" << '\n';
 
-    generate_plantuml_directives(ostr, m_config.puml().before);
+    if (m_config.puml)
+        generate_plantuml_directives(ostr, m_config.puml().before);
 
     // Generate files and folders
     util::for_each_if(
         m_model,
         [this](const auto &f) {
-            return f->type() == common::model::source_file_t::kDirectory ||
-                m_model.should_include(*f);
+            return true;
         },
         [this, &ostr](const auto &f) {
             generate(dynamic_cast<source_file &>(*f), ostr);
@@ -115,7 +120,8 @@ void generator::generate(std::ostream &ostr) const
 
     generate_config_layout_hints(ostr);
 
-    generate_plantuml_directives(ostr, m_config.puml().after);
+    if (m_config.puml)
+        generate_plantuml_directives(ostr, m_config.puml().after);
 
     ostr << "@enduml" << '\n';
 }
diff --git a/src/include_diagram/model/diagram.cc b/src/include_diagram/model/diagram.cc
index 163e330f..aceb89bb 100644
--- a/src/include_diagram/model/diagram.cc
+++ b/src/include_diagram/model/diagram.cc
@@ -35,33 +35,50 @@ common::optional_ref<common::model::diagram_element> diagram::get(
 }
 
 common::optional_ref<common::model::diagram_element> diagram::get(
-    const common::model::diagram_element::id_t /*id*/) const
+    const common::model::diagram_element::id_t id) const
 {
-    return {};
+    return get_file(id);
 }
 
 void diagram::add_file(std::unique_ptr<common::model::source_file> &&f)
 {
-    LOG_DBG("Adding source file: {}, {}", f->name(), f->full_name(true));
+    // Don't add the same file more than once
+    if(get_file(f->id()))
+        return;
 
-    files_.emplace_back(*f);
+    LOG_DBG("Adding source file: {}, {}", f->name(), f->fs_path().string());
 
-    auto p = f->path();
+    auto &ff = *f;
+
+    assert(!ff.name().empty());
+    assert(ff.id() != 0);
+
+    files_.emplace_back(ff);
+
+    auto p = ff.path();
 
     if (!f->path().is_empty()) {
         // If the parent path is not empty, ensure relative parent directories
         // of this source_file are in the diagram
         common::model::filesystem_path parent_path_so_far;
         for (const auto &directory : f->path()) {
-            auto dir = std::make_unique<common::model::source_file>();
-            if (!parent_path_so_far.is_empty())
-                dir->set_path(parent_path_so_far);
-            dir->set_name(directory);
+            auto source_file_path = parent_path_so_far | directory;
+            if(parent_path_so_far.is_empty())
+                source_file_path = {directory};
+
+            auto dir = std::make_unique<common::model::source_file>(
+                std::filesystem::path{source_file_path.to_string()});
             dir->set_type(common::model::source_file_t::kDirectory);
 
-            if (!get_element(parent_path_so_far | directory).has_value())
+            assert(!dir->name().empty());
+
+            if (!get_element(source_file_path).has_value()) {
                 add_file(std::move(dir));
 
+                LOG_DBG("Added directory '{}' at path '{}'", directory,
+                    parent_path_so_far.to_string());
+            }
+
             parent_path_so_far.append(directory);
         }
     }
@@ -81,6 +98,18 @@ common::optional_ref<common::model::source_file> diagram::get_file(
     return {};
 }
 
+common::optional_ref<common::model::source_file> diagram::get_file(
+    const common::model::diagram_element::id_t id) const
+{
+    for (const auto &p : files_) {
+        if (p.get().id() == id) {
+            return {p};
+        }
+    }
+
+    return {};
+}
+
 std::string diagram::to_alias(const std::string &full_name) const
 {
     LOG_DBG("Looking for alias for {}", full_name);
diff --git a/src/include_diagram/model/diagram.h b/src/include_diagram/model/diagram.h
index 19e9f796..11dda724 100644
--- a/src/include_diagram/model/diagram.h
+++ b/src/include_diagram/model/diagram.h
@@ -52,6 +52,9 @@ public:
     common::optional_ref<common::model::source_file> get_file(
         const std::string &name) const;
 
+    common::optional_ref<common::model::source_file> get_file(
+        const common::model::diagram_element::id_t id) const;
+
     std::string to_alias(const std::string &full_name) const;
 
     const common::reference_vector<common::model::source_file> &files() const;
diff --git a/src/include_diagram/visitor/translation_unit_visitor.cc b/src/include_diagram/visitor/translation_unit_visitor.cc
index 896ba6ea..1dc6a628 100644
--- a/src/include_diagram/visitor/translation_unit_visitor.cc
+++ b/src/include_diagram/visitor/translation_unit_visitor.cc
@@ -18,6 +18,8 @@
 
 #include "translation_unit_visitor.h"
 
+#include "common/clang_utils.h"
+
 #include <filesystem>
 
 namespace clanguml::include_diagram::visitor {
@@ -31,4 +33,203 @@ translation_unit_visitor::translation_unit_visitor(clang::SourceManager &sm,
 {
 }
 
+translation_unit_visitor::include_visitor::include_visitor(
+    clang::SourceManager &sm,
+    clanguml::include_diagram::model::diagram &diagram,
+    const clanguml::config::include_diagram &config)
+    : source_manager_{sm}
+    , diagram_{diagram}
+    , config_{config}
+{
+}
+
+void translation_unit_visitor::include_visitor::InclusionDirective(
+    clang::SourceLocation hash_loc, const clang::Token &include_tok,
+    clang::StringRef file_name, bool is_angled,
+    clang::CharSourceRange filename_range, const clang::FileEntry *file,
+    clang::StringRef search_path, clang::StringRef relative_path,
+    const clang::Module *imported, clang::SrcMgr::CharacteristicKind file_type)
+{
+    using common::model::relationship;
+    using common::model::source_file;
+    using common::model::source_file_t;
+
+    auto current_file =
+        std::filesystem::path{source_manager_.getFilename(hash_loc).str()};
+    current_file = std::filesystem::absolute(current_file);
+    current_file = current_file.lexically_normal();
+
+    auto current_file_id = process_source_file(current_file);
+    if (!current_file_id)
+        return;
+
+    assert(diagram().get(current_file_id.value()));
+
+    auto include_path = std::filesystem::path(file->getDir()->getName().str());
+    include_path = include_path / file->getName().str();
+    include_path = include_path.lexically_normal();
+
+    LOG_DBG("Processing include file {} in file {}", include_path.string(),
+        current_file.string());
+
+    auto relative_include_path = include_path;
+    if (config().relative_to) {
+        const std::filesystem::path relative_to{config().relative_to()};
+        relative_include_path =
+            std::filesystem::relative(include_path, relative_to);
+    }
+
+    if (diagram().should_include(source_file{include_path})) {
+        process_internal_header(include_path,
+            file_type != clang::SrcMgr::CharacteristicKind::C_User,
+            current_file_id.value());
+    }
+    else if (config().generate_system_headers() && is_angled) {
+        process_external_system_header(
+            relative_path.str(), current_file_id.value());
+    }
+    else {
+        LOG_DBG("Skipping include directive to file {}", include_path.string());
+    }
+}
+
+std::optional<common::id_t>
+translation_unit_visitor::include_visitor::process_internal_header(
+    const std::filesystem::path &include_path, bool is_system,
+    const common::id_t current_file_id)
+{
+    // Relativize the path with respect to relative_to config option
+    auto relative_include_path = include_path;
+    if (config().relative_to) {
+        const std::filesystem::path relative_to{config().relative_to()};
+        relative_include_path =
+            std::filesystem::relative(include_path, relative_to);
+    }
+
+    // Check if this source file is already registered in the diagram,
+    // if not add it
+    auto diagram_path =
+        common::model::source_file{relative_include_path}.full_path();
+    if (!diagram().get_element(diagram_path.to_string()).has_value()) {
+        diagram().add_file(std::make_unique<common::model::source_file>(
+            diagram_path.to_string()));
+    }
+
+    auto &include_file = diagram().get_element(diagram_path).value();
+
+    include_file.set_type(common::model::source_file_t::kHeader);
+    include_file.set_file(
+        std::filesystem::absolute(include_path).lexically_normal().string());
+    include_file.set_line(0);
+
+    // Add relationship from the currently parsed source file to this
+    // include file
+    const auto relationship_type = is_system
+        ? common::model::relationship_t::kDependency
+        : common::model::relationship_t::kAssociation;
+
+    if (diagram().get(current_file_id)) {
+        diagram()
+            .get(current_file_id)
+            .value()
+            .get()
+            .add_relationship(common::model::relationship{
+                relationship_type, include_file.id()});
+    }
+
+    return include_file.id();
+}
+
+std::optional<common::id_t>
+translation_unit_visitor::include_visitor::process_external_system_header(
+    const std::filesystem::path &include_path,
+    const common::id_t current_file_id)
+{
+    const auto file_name = include_path.filename();
+    const auto file_name_str = file_name.string();
+
+    auto f = std::make_unique<common::model::source_file>();
+    f->set_name(include_path.string());
+    f->set_type(common::model::source_file_t::kHeader);
+    f->set_id(common::to_id(include_path));
+
+    const auto f_id = f->id();
+
+    diagram().add_file(std::move(f));
+
+    if (diagram().get(current_file_id)) {
+        diagram()
+            .get(current_file_id)
+            .value()
+            .get()
+            .add_relationship(common::model::relationship{
+                common::model::relationship_t::kDependency, f_id});
+    }
+
+    return f_id;
+}
+
+std::optional<common::id_t>
+translation_unit_visitor::include_visitor::process_source_file(
+    const std::filesystem::path &file)
+{
+    using common::model::relationship;
+    using common::model::source_file;
+    using common::model::source_file_t;
+
+    LOG_DBG("Processing source file {}", file.string());
+
+    auto file_path = std::filesystem::path{file};
+
+    // Make sure the file_path is absolute with respect to the
+    // filesystem, and in normal form
+    if (file_path.is_relative()) {
+        file_path = config().base_directory() / file_path;
+    }
+
+    file_path = file_path.lexically_normal();
+
+    // Relativize the path with respect to relative_to config option
+    auto relative_file_path = file_path;
+    if (config().relative_to) {
+        const std::filesystem::path relative_to{config().relative_to()};
+        relative_file_path = std::filesystem::relative(file_path, relative_to);
+    }
+
+    if (diagram().should_include(source_file{file_path})) {
+        [[maybe_unused]] const auto relative_file_path_str =
+            relative_file_path.string();
+
+        // Check if this source file is already registered in the diagram,
+        // if not add it
+        auto diagram_path = source_file{relative_file_path}.full_path();
+        if (!diagram().get_element(diagram_path).has_value()) {
+            diagram().add_file(
+                std::make_unique<source_file>(relative_file_path));
+        }
+
+        auto &source_file = diagram().get_element(diagram_path).value();
+
+        const std::string implementation_suffix_prefix{".c"};
+        if (file_path.has_extension() &&
+            util::starts_with(
+                file_path.extension().string(), implementation_suffix_prefix)) {
+            source_file.set_type(source_file_t::kImplementation);
+        }
+        else
+            source_file.set_type(source_file_t::kHeader);
+
+        source_file.set_file(std::filesystem::absolute(file.string())
+                                 .lexically_normal()
+                                 .string());
+        source_file.set_line(0);
+
+        return source_file.id();
+    }
+
+    LOG_DBG("Skipping source file {}", file_path.string());
+
+    return {};
+}
+
 }
diff --git a/src/include_diagram/visitor/translation_unit_visitor.h b/src/include_diagram/visitor/translation_unit_visitor.h
index a587de99..0b3a65bb 100644
--- a/src/include_diagram/visitor/translation_unit_visitor.h
+++ b/src/include_diagram/visitor/translation_unit_visitor.h
@@ -25,6 +25,7 @@
 
 #include <clang/AST/RecursiveASTVisitor.h>
 #include <clang/Basic/SourceManager.h>
+#include <clang/Lex/PPCallbacks.h>
 
 #include <functional>
 #include <map>
@@ -36,11 +37,57 @@ namespace clanguml::include_diagram::visitor {
 class translation_unit_visitor
     : public clang::RecursiveASTVisitor<translation_unit_visitor> {
 public:
+    // This is an internal class for convenience to be able to access the
+    // include_visitor type from translation_unit_visitor type
+    class include_visitor : public clang::PPCallbacks {
+    public:
+        include_visitor(clang::SourceManager &sm,
+            clanguml::include_diagram::model::diagram &diagram,
+            const clanguml::config::include_diagram &config);
+
+        void InclusionDirective(clang::SourceLocation hash_loc,
+            const clang::Token &include_tok, clang::StringRef file_name,
+            bool is_angled, clang::CharSourceRange filename_range,
+            const clang::FileEntry *file, clang::StringRef search_path,
+            clang::StringRef relative_path, const clang::Module *imported,
+            clang::SrcMgr::CharacteristicKind file_type) override;
+
+        std::optional<common::id_t> process_internal_header(const std::filesystem::path &include_path,
+            bool is_system, const common::id_t current_file_id);
+
+        std::optional<common::id_t> process_external_system_header(
+            const std::filesystem::path &include_path,
+            const common::id_t current_file_id);
+
+        std::optional<common::id_t> process_source_file(const std::filesystem::path &file);
+
+        clanguml::include_diagram::model::diagram &diagram()
+        {
+            return diagram_;
+        }
+
+        const clanguml::config::include_diagram &config() const
+        {
+            return config_;
+        }
+
+    private:
+        clang::SourceManager &source_manager_;
+
+        // Reference to the output diagram model
+        clanguml::include_diagram::model::diagram &diagram_;
+
+        // Reference to class diagram config
+        const clanguml::config::include_diagram &config_;
+    };
+
     translation_unit_visitor(clang::SourceManager &sm,
         clanguml::include_diagram::model::diagram &diagram,
         const clanguml::config::include_diagram &config);
 
-    void operator()(const cppast::cpp_entity &file);
+    clanguml::include_diagram::model::diagram &diagram() { return diagram_; }
+
+    const clanguml::config::include_diagram &config() const { return config_; }
 
     void finalize() { }
 
diff --git a/tests/t40001/.clang-uml b/tests/t40001/.clang-uml
index dc6d2596..918e7f64 100644
--- a/tests/t40001/.clang-uml
+++ b/tests/t40001/.clang-uml
@@ -9,13 +9,13 @@ diagrams:
       - ../../tests/t40001/**/*.cc
       - ../../tests/t40001/**/*.h
     # Render the paths relative to this directory
-    relative_to: ../../tests/t40001
+    relative_to: ../../../tests/t40001
     # Include also external system headers
     generate_system_headers: true
     include:
       # Include only headers belonging to these paths
       paths:
-        - ../../tests/t40001
+        - ../../../tests/t40001
     plantuml:
       before:
         - "' t40001 test include diagram"
diff --git a/tests/t40001/test_case.h b/tests/t40001/test_case.h
index 5f476c27..2f089af5 100644
--- a/tests/t40001/test_case.h
+++ b/tests/t40001/test_case.h
@@ -24,7 +24,7 @@ TEST_CASE("t40001", "[test-case][package]")
 
     REQUIRE(diagram->name == "t40001_include");
 
-    auto model = generate_include_diagram(db, diagram);
+    auto model = generate_include_diagram(*db, diagram);
 
     REQUIRE(model->name() == "t40001_include");
 
diff --git a/tests/t40002/.clang-uml b/tests/t40002/.clang-uml
index 69813945..8daf8264 100644
--- a/tests/t40002/.clang-uml
+++ b/tests/t40002/.clang-uml
@@ -9,15 +9,15 @@ diagrams:
       - ../../tests/t40002/**/*.cc
       - ../../tests/t40002/**/*.h
     # Render the paths relative to this directory
-    relative_to: ../../tests/t40002
+    relative_to: ../../../tests/t40002
     include:
       # Include only files belonging to these paths
       paths:
-        - ../../tests/t40002
+        - ../../../tests/t40002
     exclude:
       paths:
         # Exclude single header
-        - ../../tests/t40002/include/lib2/lib2_detail.h
+        - ../../../tests/t40002/include/lib2/lib2_detail.h
     plantuml:
       before:
         - "' t40002 test include diagram"
\ No newline at end of file
diff --git a/tests/t40002/test_case.h b/tests/t40002/test_case.h
index 1baefec5..15aa503a 100644
--- a/tests/t40002/test_case.h
+++ b/tests/t40002/test_case.h
@@ -24,7 +24,7 @@ TEST_CASE("t40002", "[test-case][package]")
 
     REQUIRE(diagram->name == "t40002_include");
 
-    auto model = generate_include_diagram(db, diagram);
+    auto model = generate_include_diagram(*db, diagram);
 
     REQUIRE(model->name() == "t40002_include");
 
diff --git a/tests/t40003/.clang-uml b/tests/t40003/.clang-uml
index df927287..64f03790 100644
--- a/tests/t40003/.clang-uml
+++ b/tests/t40003/.clang-uml
@@ -9,13 +9,14 @@ diagrams:
       - ../../tests/t40003/include/**/*.h
       - ../../tests/t40003/src/**/*.cc
     # Render the paths relative to this directory
-    relative_to: ../../tests/t40003
+    relative_to: ../../../tests/t40003
     include:
-      # Include only files belonging to these paths
+      # Include only files which depend on t1.h
       dependants:
-        - ../../tests/t40003/include/dependants/t1.h
+        - ../../../tests/t40003/include/dependants/t1.h
+      # and dependencies of t2.cc
       dependencies:
-        - ../../tests/t40003/src/dependencies/t2.cc
+        - ../../../tests/t40003/src/dependencies/t2.cc
     plantuml:
       before:
         - "' t40003 test include diagram"
\ No newline at end of file
diff --git a/tests/t40003/test_case.h b/tests/t40003/test_case.h
index 38c2c3fd..9ad79ce2 100644
--- a/tests/t40003/test_case.h
+++ b/tests/t40003/test_case.h
@@ -24,7 +24,7 @@ TEST_CASE("t40003", "[test-case][package]")
 
     REQUIRE(diagram->name == "t40003_include");
 
-    auto model = generate_include_diagram(db, diagram);
+    auto model = generate_include_diagram(*db, diagram);
 
     REQUIRE(model->name() == "t40003_include");
 
diff --git a/tests/test_cases.cc b/tests/test_cases.cc
index 7f98c0c9..8e5fe618 100644
--- a/tests/test_cases.cc
+++ b/tests/test_cases.cc
@@ -257,9 +257,9 @@ using namespace clanguml::test::matchers;
 ////
 //// Include diagram tests
 ////
-//#include "t40001/test_case.h"
-//#include "t40002/test_case.h"
-//#include "t40003/test_case.h"
+#include "t40001/test_case.h"
+#include "t40002/test_case.h"
+#include "t40003/test_case.h"
 //
 ////
 //// Other tests (e.g. configuration file)