From 17b48211968f82ca36b9b416be4e61828fbd03b6 Mon Sep 17 00:00:00 2001
From: LProcopi15 <leah.procopi@gmail.com>
Date: Thu, 16 Jan 2025 17:22:08 -0800
Subject: [PATCH] add unit tests & make final edits to functional tests

---
 tests/functional/iceberg/models.py  |  8 +--
 tests/unit/test_iceberg_location.py | 81 +++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 4 deletions(-)
 create mode 100644 tests/unit/test_iceberg_location.py

diff --git a/tests/functional/iceberg/models.py b/tests/functional/iceberg/models.py
index 614cb4a79..606792b36 100644
--- a/tests/functional/iceberg/models.py
+++ b/tests/functional/iceberg/models.py
@@ -31,7 +31,7 @@
     cluster_by=['id'],
     table_format="iceberg",
     external_volume="s3_iceberg_snow",
-    base_location_root="base_path",
+    base_location_root="root_path",
   )
 }}
 
@@ -46,7 +46,7 @@
     cluster_by=['id'],
     table_format="iceberg",
     external_volume="s3_iceberg_snow",
-    base_location_root="base_path",
+    base_location_root="root_path",
     base_location_subpath="subpath",
   )
 }}
@@ -75,7 +75,7 @@
     cluster_by=['id'],
     table_format="iceberg",
     external_volume="s3_iceberg_snow",
-    base_location_root="base_path",
+    base_location_root="root_path",
   )
 }}
 
@@ -90,7 +90,7 @@
     cluster_by=['id'],
     table_format="iceberg",
     external_volume="s3_iceberg_snow",
-    base_location_root="base_path",
+    base_location_root="root_path",
     base_location_subpath='subpath',
   )
 }}
diff --git a/tests/unit/test_iceberg_location.py b/tests/unit/test_iceberg_location.py
new file mode 100644
index 000000000..58529e4e9
--- /dev/null
+++ b/tests/unit/test_iceberg_location.py
@@ -0,0 +1,81 @@
+import pytest
+from dbt.adapters.snowflake.relation import SnowflakeRelation
+
+
+@pytest.fixture
+def iceberg_config() -> dict:
+    """Fixture providing standard Iceberg configuration."""
+    return {
+        "schema": "my_schema",
+        "identifier": "my_table",
+        "external_volume": "s3_iceberg_snow",
+        "base_location_root": "root_path",
+        "base_location_subpath": "subpath",
+    }
+
+
+def get_actual_base_location(config: dict[str, str]) -> str:
+    """Get the actual base location from the configuration by parsing the DDL predicates."""
+
+    relation = SnowflakeRelation.create(
+        schema=config["schema"],
+        identifier=config["identifier"],
+    )
+
+    actual_ddl_predicates = relation.get_iceberg_ddl_options(config).strip()
+    actual_base_location = actual_ddl_predicates.split("base_location = ")[1]
+
+    return actual_base_location
+
+
+def test_iceberg_path_and_subpath(iceberg_config: dict[str, str]):
+    """Test when base_location_root and base_location_subpath are provided"""
+    expected_base_location = (
+        f"'{iceberg_config['base_location_root']}/"
+        f"{iceberg_config['schema']}/"
+        f"{iceberg_config['identifier']}/"
+        f"{iceberg_config['base_location_subpath']}'"
+    ).strip()
+
+    assert get_actual_base_location(iceberg_config) == expected_base_location
+
+
+def test_iceberg_only_subpath(iceberg_config: dict[str, str]):
+    """Test when only base_location_subpath is provided"""
+    del iceberg_config["base_location_root"]
+
+    expected_base_location = (
+        f"'_dbt/"
+        f"{iceberg_config['schema']}/"
+        f"{iceberg_config['identifier']}/"
+        f"{iceberg_config['base_location_subpath']}'"
+    ).strip()
+
+    assert get_actual_base_location(iceberg_config) == expected_base_location
+
+
+def test_iceberg_only_path(iceberg_config: dict[str, str]):
+    """Test when only base_location_root is provided"""
+    del iceberg_config["base_location_subpath"]
+
+    expected_base_location = (
+        f"'{iceberg_config['base_location_root']}/"
+        f"{iceberg_config['schema']}/"
+        f"{iceberg_config['identifier']}'"
+    ).strip()
+
+    assert get_actual_base_location(iceberg_config) == expected_base_location
+
+
+def test_iceberg_no_path(iceberg_config: dict[str, str]):
+    """Test when no base_location_root or is base_location_subpath provided"""
+    del iceberg_config["base_location_root"]
+    del iceberg_config["base_location_subpath"]
+
+    expected_base_location = (
+        f"'_dbt/"
+        f"{iceberg_config['schema']}/"
+        f"{iceberg_config['identifier']}'"
+    ).strip()
+
+    assert get_actual_base_location(iceberg_config) == expected_base_location