diff --git a/src/dbt_osmosis/core/osmosis.py b/src/dbt_osmosis/core/osmosis.py
index 364866d..a663908 100644
--- a/src/dbt_osmosis/core/osmosis.py
+++ b/src/dbt_osmosis/core/osmosis.py
@@ -6,7 +6,6 @@
 import json
 import os
 import re
-import sys
 import threading
 import time
 import typing as t
@@ -15,7 +14,7 @@
 from collections.abc import Iterable, Iterator
 from concurrent.futures import FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait
 from dataclasses import dataclass, field
-from functools import lru_cache, partial
+from functools import lru_cache
 from itertools import chain
 from pathlib import Path
 
@@ -547,7 +546,7 @@ def normalize_column_name(column: str, credentials_type: str) -> str:
     return column
 
 
-def _maybe_use_precise_dtype(col: t.Any, settings: YamlRefactorSettings) -> str:
+def _maybe_use_precise_dtype(col: BaseColumn, settings: YamlRefactorSettings) -> str:
     """Use the precise data type if enabled in the settings."""
     if (col.is_numeric() and settings.numeric_precision) or (
         col.is_string() and settings.char_length
diff --git a/tests/test_yaml_context.py b/tests/test_yaml_context.py
new file mode 100644
index 0000000..857918e
--- /dev/null
+++ b/tests/test_yaml_context.py
@@ -0,0 +1,126 @@
+from unittest import mock
+
+import pytest
+
+from dbt_osmosis.core.osmosis import (
+    DbtConfiguration,
+    YamlRefactorContext,
+    YamlRefactorSettings,
+    apply_restructure_plan,
+    create_dbt_project_context,
+    create_missing_source_yamls,
+    draft_restructure_delta_plan,
+    get_columns,
+    get_table_ref,
+    inherit_upstream_column_knowledge,
+    reload_manifest,
+)
+
+
+@pytest.fixture(scope="module")
+def yaml_context() -> YamlRefactorContext:
+    # initializing the context is a sanity test in and of itself
+    c = DbtConfiguration(project_dir="demo_duckdb", profiles_dir="demo_duckdb")
+    c.vars = {"dbt-osmosis": {}}
+    project = create_dbt_project_context(c)
+    context = YamlRefactorContext(
+        project, settings=YamlRefactorSettings(use_unrendered_descriptions=True, dry_run=True)
+    )
+    return context
+
+
+# Sanity tests
+
+
+def test_reload_manifest(yaml_context: YamlRefactorContext):
+    reload_manifest(yaml_context.project)
+
+
+def test_create_missing_source_yamls(yaml_context: YamlRefactorContext):
+    create_missing_source_yamls(yaml_context)
+
+
+def test_draft_restructure_delta_plan(yaml_context: YamlRefactorContext):
+    assert draft_restructure_delta_plan(yaml_context) is not None
+
+
+def test_apply_restructure_plan(yaml_context: YamlRefactorContext):
+    plan = draft_restructure_delta_plan(yaml_context)
+    apply_restructure_plan(yaml_context, plan, confirm=False)
+
+
+def test_inherit_upstream_column_knowledge(yaml_context: YamlRefactorContext):
+    inherit_upstream_column_knowledge(yaml_context)
+
+
+# Column type + settings tests
+
+
+def _customer_column_types(yaml_context: YamlRefactorContext) -> dict[str, str]:
+    node = next(n for n in yaml_context.project.manifest.nodes.values() if n.name == "customers")
+    assert node
+
+    catalog_key = get_table_ref(node)
+    columns = get_columns(yaml_context, catalog_key)
+    assert columns
+
+    column_types = dict({name: meta.type for name, meta in columns.items()})
+    assert column_types
+    return column_types
+
+
+def test_get_columns_meta(yaml_context: YamlRefactorContext):
+    with mock.patch("dbt_osmosis.core.osmosis._COLUMN_LIST_CACHE", {}):
+        assert _customer_column_types(yaml_context) == {
+            # in DuckDB decimals always have presision and scale
+            "customer_average_value": "DECIMAL(18,3)",
+            "customer_id": "INTEGER",
+            "customer_lifetime_value": "DOUBLE",
+            "first_name": "VARCHAR",
+            "first_order": "DATE",
+            "last_name": "VARCHAR",
+            "most_recent_order": "DATE",
+            "number_of_orders": "BIGINT",
+        }
+
+
+def test_get_columns_meta_char_length():
+    yaml_context = YamlRefactorContext(
+        project=create_dbt_project_context(
+            DbtConfiguration(project_dir="demo_duckdb", profiles_dir="demo_duckdb")
+        ),
+        settings=YamlRefactorSettings(char_length=True, dry_run=True),
+    )
+    with mock.patch("dbt_osmosis.core.osmosis._COLUMN_LIST_CACHE", {}):
+        assert _customer_column_types(yaml_context) == {
+            # in DuckDB decimals always have presision and scale
+            "customer_average_value": "DECIMAL(18,3)",
+            "customer_id": "INTEGER",
+            "customer_lifetime_value": "DOUBLE",
+            "first_name": "character varying(256)",
+            "first_order": "DATE",
+            "last_name": "character varying(256)",
+            "most_recent_order": "DATE",
+            "number_of_orders": "BIGINT",
+        }
+
+
+def test_get_columns_meta_numeric_precision():
+    yaml_context = YamlRefactorContext(
+        project=create_dbt_project_context(
+            DbtConfiguration(project_dir="demo_duckdb", profiles_dir="demo_duckdb")
+        ),
+        settings=YamlRefactorSettings(numeric_precision=True, dry_run=True),
+    )
+    with mock.patch("dbt_osmosis.core.osmosis._COLUMN_LIST_CACHE", {}):
+        assert _customer_column_types(yaml_context) == {
+            # in DuckDB decimals always have presision and scale
+            "customer_average_value": "DECIMAL(18,3)",
+            "customer_id": "INTEGER",
+            "customer_lifetime_value": "DOUBLE",
+            "first_name": "VARCHAR",
+            "first_order": "DATE",
+            "last_name": "VARCHAR",
+            "most_recent_order": "DATE",
+            "number_of_orders": "BIGINT",
+        }
diff --git a/tests/test_yaml_manager.py b/tests/test_yaml_manager.py
deleted file mode 100644
index 024dfd4..0000000
--- a/tests/test_yaml_manager.py
+++ /dev/null
@@ -1,117 +0,0 @@
-# TODO: refactor this test
-# from pathlib import Path
-#
-# import pytest
-# from dbt.contracts.results import CatalogKey
-#
-# from dbt_osmosis.core.osmosis import DbtYamlManager
-#
-#
-# @pytest.fixture(scope="module")
-# def yaml_manager() -> DbtYamlManager:
-#     return DbtYamlManager(project_dir="demo_duckdb", profiles_dir="demo_duckdb", dry_run=True)
-#
-#
-# def test_initialize_adapter(yaml_manager: DbtYamlManager):
-#     yaml_manager.initialize_adapter()
-#
-#
-# def test_list(yaml_manager: DbtYamlManager):
-#     yaml_manager.list()
-#
-#
-# def test_test(yaml_manager: DbtYamlManager):
-#     yaml_manager.test()
-#
-#
-# def test_run(yaml_manager: DbtYamlManager):
-#     yaml_manager.run()
-#
-#
-# def test_build(yaml_manager: DbtYamlManager):
-#     yaml_manager.build()
-#
-#
-# def test_parse_project(yaml_manager: DbtYamlManager):
-#     yaml_manager.parse_project()
-#
-#
-# def test_safe_parse_project(yaml_manager: DbtYamlManager):
-#     yaml_manager.safe_parse_project()
-#
-#
-# def test_bootstrap_sources(yaml_manager: DbtYamlManager):
-#     yaml_manager.bootstrap_sources()
-#
-#
-# def test_draft_project_structure_update_plan(yaml_manager: DbtYamlManager):
-#     yaml_manager.draft_project_structure_update_plan()
-#
-#
-# def test_commit_project_restructure_to_disk(yaml_manager: DbtYamlManager):
-#     yaml_manager.commit_project_restructure_to_disk()
-#
-#
-# def test_propagate_documentation_downstream(yaml_manager: DbtYamlManager):
-#     yaml_manager.propagate_documentation_downstream()
-#
-#
-# def _customer_column_types(yaml_manager: DbtYamlManager) -> dict[str, str]:
-#     node = next(n for n in yaml_manager.manifest.nodes.values() if n.name == "customers")
-#     assert node
-#
-#     catalog_key = yaml_manager.get_catalog_key(node)
-#     columns = yaml_manager.get_columns_meta(catalog_key)
-#     assert columns
-#
-#     column_types = dict({name: meta.type for name, meta in columns.items()})
-#     assert column_types
-#     return column_types
-#
-#
-# def test_get_columns_meta(yaml_manager: DbtYamlManager):
-#     assert _customer_column_types(yaml_manager) == {
-#         # in DuckDB decimals always have presision and scale
-#         "customer_average_value": "DECIMAL(18,3)",
-#         "customer_id": "INTEGER",
-#         "customer_lifetime_value": "DOUBLE",
-#         "first_name": "VARCHAR",
-#         "first_order": "DATE",
-#         "last_name": "VARCHAR",
-#         "most_recent_order": "DATE",
-#         "number_of_orders": "BIGINT",
-#     }
-#
-#
-# def test_get_columns_meta_char_length():
-#     yaml_manager = DbtYamlManager(
-#         project_dir="demo_duckdb", profiles_dir="demo_duckdb", char_length=True, dry_run=True
-#     )
-#     assert _customer_column_types(yaml_manager) == {
-#         # in DuckDB decimals always have presision and scale
-#         "customer_average_value": "DECIMAL(18,3)",
-#         "customer_id": "INTEGER",
-#         "customer_lifetime_value": "DOUBLE",
-#         "first_name": "character varying(256)",
-#         "first_order": "DATE",
-#         "last_name": "character varying(256)",
-#         "most_recent_order": "DATE",
-#         "number_of_orders": "BIGINT",
-#     }
-#
-#
-# def test_get_columns_meta_numeric_precision():
-#     yaml_manager = DbtYamlManager(
-#         project_dir="demo_duckdb", profiles_dir="demo_duckdb", numeric_precision=True, dry_run=True
-#     )
-#     assert _customer_column_types(yaml_manager) == {
-#         # in DuckDB decimals always have presision and scale
-#         "customer_average_value": "DECIMAL(18,3)",
-#         "customer_id": "INTEGER",
-#         "customer_lifetime_value": "DOUBLE",
-#         "first_name": "VARCHAR",
-#         "first_order": "DATE",
-#         "last_name": "VARCHAR",
-#         "most_recent_order": "DATE",
-#         "number_of_orders": "BIGINT",
-#     }