From f9e9917a8c6def12f0ac1627ed47b275d51e8d13 Mon Sep 17 00:00:00 2001
From: Zane Clark <zane.clark@pwc.com>
Date: Thu, 31 Oct 2024 07:48:46 -0700
Subject: [PATCH] feat: swap tomllib with tomlkit

---
 tests/config/test_get_merged_config.py | 4 ++--
 tests/test_main.py                     | 6 ++++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tests/config/test_get_merged_config.py b/tests/config/test_get_merged_config.py
index 8725b9b..9ada158 100644
--- a/tests/config/test_get_merged_config.py
+++ b/tests/config/test_get_merged_config.py
@@ -1,7 +1,7 @@
 import json
 import logging
 import os
-import tomllib
+import tomlkit
 from pathlib import Path
 from unittest import mock
 
@@ -23,7 +23,7 @@
 
 def get_connection_from_toml(file_path: Path, connection_name: str) -> dict:
     with file_path.open("rb") as f:
-        connections = tomllib.load(f)
+        connections = tomlkit.load(f)
         return connections[connection_name]
 
 
diff --git a/tests/test_main.py b/tests/test_main.py
index ae78874..178f2aa 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -3,7 +3,7 @@
 import json
 import logging
 import os
-import tomllib
+import tomlkit
 import tempfile
 import unittest.mock as mock
 from dataclasses import asdict
@@ -21,7 +21,7 @@
 
 def get_connection_from_toml(file_path: Path, connection_name: str) -> dict:
     with file_path.open("rb") as f:
-        connections = tomllib.load(f)
+        connections = tomlkit.load(f)
         return connections[connection_name]
 
 
@@ -603,6 +603,7 @@ def test_main_deploy_config_folder(
                     )
                 )
 
+            # noinspection PyTypeChecker
             args[args.index("DUMMY")] = d
             expected_config["config_file_path"] = Path(d) / "schemachange-config.yml"
 
@@ -654,6 +655,7 @@ def test_main_deploy_modules_folder(
 ):
     with mock.patch.dict(os.environ, {"SNOWFLAKE_PASSWORD": "password"}, clear=True):
         with tempfile.TemporaryDirectory() as d:
+            # noinspection PyTypeChecker
             args[args.index("DUMMY")] = d
             expected_config["modules_folder"] = Path(d)