test(python): add tests with mock APIs for Unity Catalog variants

omkar-foss · ion-elgreco · commit 3445b337c0da · 2025-02-18T20:41:22.000Z
Signed-off-by: Omkar P &lt;45419097+omkar-foss@users.noreply.github.com&gt;
diff --git a/.github/workflows/mockoon_data_files/unitycatalog_databricks.json b/.github/workflows/mockoon_data_files/unitycatalog_databricks.json
@@ -0,0 +1,112 @@
+{
+  "uuid": "0696d6d5-62b6-4a0f-8524-917b0c3848d4",
+  "lastMigration": 33,
+  "name": "Unity Catalog (Databricks) Test API",
+  "endpointPrefix": "",
+  "latency": 0,
+  "port": 8080,
+  "hostname": "",
+  "folders": [],
+  "routes": [
+    {
+      "uuid": "16d88cdc-abcf-4b0d-a4d6-18a3205a9a6c",
+      "type": "http",
+      "documentation": "Temporary Table Credentials API",
+      "method": "post",
+      "endpoint": "api/2.1/unity-catalog/temporary-table-credentials",
+      "responses": [
+        {
+          "uuid": "50bba5c1-1e2c-4d5b-a611-ab478320f2b0",
+          "body": "{\n  \"aws_temp_credentials\": {\n    \"access_key_id\": \"string\",\n    \"secret_access_key\": \"string\",\n    \"session_token\": \"string\",\n    \"access_point\": \"string\"\n  },\n  \"azure_user_delegation_sas\": {\n    \"sas_token\": \"string\"\n  },\n  \"r2_temp_credentials\": {\n    \"access_key_id\": \"string\",\n    \"secret_access_key\": \"string\",\n    \"session_token\": \"string\"\n  },\n  \"expiration_time\": 0,\n  \"url\": \"string\"\n}",
+          "latency": 0,
+          "statusCode": 200,
+          "label": "Default response",
+          "headers": [],
+          "bodyType": "INLINE",
+          "filePath": "",
+          "databucketID": "",
+          "sendFileAsBody": false,
+          "rules": [],
+          "rulesOperator": "OR",
+          "disableTemplating": false,
+          "fallbackTo404": false,
+          "default": true,
+          "crudKey": "id",
+          "callbacks": []
+        }
+      ],
+      "responseMode": null,
+      "streamingMode": null,
+      "streamingInterval": 0
+    },
+    {
+      "uuid": "66f51ad5-2ba3-4cb0-b633-c362a9f75836",
+      "type": "http",
+      "documentation": "Get Table Details API",
+      "method": "get",
+      "endpoint": "api/2.1/unity-catalog/tables/unity.default.testtable",
+      "responses": [
+        {
+          "uuid": "da280b2c-07b6-45ee-8da7-92c50e4e540e",
+          "body": "{\n  \"name\": \"string\",\n  \"catalog_name\": \"string\",\n  \"schema_name\": \"string\",\n  \"table_type\": \"MANAGED\",\n  \"data_source_format\": \"DELTA\",\n  \"columns\": [\n    {\n      \"name\": \"string\",\n      \"type_text\": \"string\",\n      \"type_name\": \"BOOLEAN\",\n      \"position\": 0,\n      \"type_precision\": 0,\n      \"type_scale\": 0,\n      \"type_interval_type\": \"string\",\n      \"type_json\": \"string\",\n      \"comment\": \"string\",\n      \"nullable\": true,\n      \"partition_index\": 0,\n      \"mask\": {\n        \"function_name\": \"string\",\n        \"using_column_names\": [\n          \"string\"\n        ]\n      }\n    }\n  ],\n  \"storage_location\": \"../crates/test/tests/data/delta-0.8.0-partitioned\",\n  \"view_definition\": \"string\",\n  \"view_dependencies\": {\n    \"dependencies\": [\n      {\n        \"table\": {\n          \"table_full_name\": \"string\"\n        },\n        \"function\": {\n          \"function_full_name\": \"string\"\n        }\n      }\n    ]\n  },\n  \"sql_path\": \"string\",\n  \"owner\": \"string\",\n  \"comment\": \"string\",\n  \"properties\": {\n    \"property1\": \"string\",\n    \"property2\": \"string\"\n  },\n  \"storage_credential_name\": \"string\",\n  \"table_constraints\": [\n    {\n      \"primary_key_constraint\": {\n        \"name\": \"string\",\n        \"child_columns\": [\n          \"string\"\n        ]\n      },\n      \"foreign_key_constraint\": {\n        \"name\": \"string\",\n        \"child_columns\": [\n          \"string\"\n        ],\n        \"parent_table\": \"string\",\n        \"parent_columns\": [\n          \"string\"\n        ]\n      },\n      \"named_table_constraint\": {\n        \"name\": \"string\"\n      }\n    }\n  ],\n  \"row_filter\": {\n    \"function_name\": \"string\",\n    \"input_column_names\": [\n      \"string\"\n    ]\n  },\n  \"enable_predictive_optimization\": \"DISABLE\",\n  \"metastore_id\": \"string\",\n  \"full_name\": \"string\",\n  \"data_access_configuration_id\": \"string\",\n  \"created_at\": 0,\n  \"created_by\": \"string\",\n  \"updated_at\": 0,\n  \"updated_by\": \"string\",\n  \"deleted_at\": 0,\n  \"table_id\": \"string\",\n  \"delta_runtime_properties_kvpairs\": {\n    \"delta_runtime_properties\": {\n      \"property1\": \"string\",\n      \"property2\": \"string\"\n    }\n  },\n  \"effective_predictive_optimization_flag\": {\n    \"value\": \"DISABLE\",\n    \"inherited_from_type\": \"CATALOG\",\n    \"inherited_from_name\": \"string\"\n  },\n  \"access_point\": \"string\",\n  \"pipeline_id\": \"string\",\n  \"browse_only\": true\n}",
+          "latency": 0,
+          "statusCode": 200,
+          "label": "Default response",
+          "headers": [],
+          "bodyType": "INLINE",
+          "filePath": "",
+          "databucketID": "",
+          "sendFileAsBody": false,
+          "rules": [],
+          "rulesOperator": "OR",
+          "disableTemplating": false,
+          "fallbackTo404": false,
+          "default": true,
+          "crudKey": "id",
+          "callbacks": []
+        }
+      ],
+      "responseMode": null,
+      "streamingMode": null,
+      "streamingInterval": 0
+    }
+  ],
+  "rootChildren": [
+    {
+      "type": "route",
+      "uuid": "16d88cdc-abcf-4b0d-a4d6-18a3205a9a6c"
+    },
+    {
+      "type": "route",
+      "uuid": "66f51ad5-2ba3-4cb0-b633-c362a9f75836"
+    }
+  ],
+  "proxyMode": false,
+  "proxyHost": "",
+  "proxyRemovePrefix": false,
+  "tlsOptions": {
+    "enabled": false,
+    "type": "CERT",
+    "pfxPath": "",
+    "certPath": "",
+    "keyPath": "",
+    "caPath": "",
+    "passphrase": ""
+  },
+  "cors": true,
+  "headers": [],
+  "proxyReqHeaders": [
+    {
+      "key": "",
+      "value": ""
+    }
+  ],
+  "proxyResHeaders": [
+    {
+      "key": "",
+      "value": ""
+    }
+  ],
+  "data": [],
+  "callbacks": []
+}
diff --git a/.github/workflows/mockoon_data_files/unitycatalog_oss.json b/.github/workflows/mockoon_data_files/unitycatalog_oss.json
@@ -0,0 +1,112 @@
+{
+  "uuid": "0696d6d5-62b6-4a0f-8524-917b0c3848d4",
+  "lastMigration": 33,
+  "name": "Unity Catalog (OSS) Test API",
+  "endpointPrefix": "",
+  "latency": 0,
+  "port": 8080,
+  "hostname": "",
+  "folders": [],
+  "routes": [
+    {
+      "uuid": "16d88cdc-abcf-4b0d-a4d6-18a3205a9a6c",
+      "type": "http",
+      "documentation": "Temporary Table Credentials API",
+      "method": "post",
+      "endpoint": "api/2.1/unity-catalog/temporary-table-credentials",
+      "responses": [
+        {
+          "uuid": "50bba5c1-1e2c-4d5b-a611-ab478320f2b0",
+          "body": "{\n  \"aws_temp_credentials\": {\n    \"access_key_id\": \"string\",\n    \"secret_access_key\": \"string\",\n    \"session_token\": \"string\"\n  },\n  \"azure_user_delegation_sas\": {\n    \"sas_token\": \"string\"\n  },\n  \"gcp_oauth_token\": {\n    \"oauth_token\": \"string\"\n  },\n  \"expiration_time\": 0\n}",
+          "latency": 0,
+          "statusCode": 200,
+          "label": "Default response",
+          "headers": [],
+          "bodyType": "INLINE",
+          "filePath": "",
+          "databucketID": "",
+          "sendFileAsBody": false,
+          "rules": [],
+          "rulesOperator": "OR",
+          "disableTemplating": false,
+          "fallbackTo404": false,
+          "default": true,
+          "crudKey": "id",
+          "callbacks": []
+        }
+      ],
+      "responseMode": null,
+      "streamingMode": null,
+      "streamingInterval": 0
+    },
+    {
+      "uuid": "66f51ad5-2ba3-4cb0-b633-c362a9f75836",
+      "type": "http",
+      "documentation": "Get Table Details API",
+      "method": "get",
+      "endpoint": "api/2.1/unity-catalog/tables/unity.default.testtable",
+      "responses": [
+        {
+          "uuid": "da280b2c-07b6-45ee-8da7-92c50e4e540e",
+          "body": "{\n  \"name\": \"string\",\n  \"catalog_name\": \"string\",\n  \"schema_name\": \"string\",\n  \"table_type\": \"MANAGED\",\n  \"data_source_format\": \"DELTA\",\n  \"columns\": [\n    {\n      \"name\": \"string\",\n      \"type_text\": \"string\",\n      \"type_json\": \"string\",\n      \"type_name\": \"BOOLEAN\",\n      \"type_precision\": 0,\n      \"type_scale\": 0,\n      \"type_interval_type\": \"string\",\n      \"position\": 0,\n      \"comment\": \"string\",\n      \"nullable\": true,\n      \"partition_index\": 0\n    }\n  ],\n  \"storage_location\": \"string\",\n  \"comment\": \"string\",\n  \"properties\": {\n    \"additionalProp1\": \"string\",\n    \"additionalProp2\": \"string\",\n    \"additionalProp3\": \"string\"\n  },\n  \"owner\": \"string\",\n  \"created_at\": 0,\n  \"created_by\": \"string\",\n  \"updated_at\": 0,\n  \"updated_by\": \"string\",\n  \"table_id\": \"string\"\n}",
+          "latency": 0,
+          "statusCode": 200,
+          "label": "Default response",
+          "headers": [],
+          "bodyType": "INLINE",
+          "filePath": "",
+          "databucketID": "",
+          "sendFileAsBody": false,
+          "rules": [],
+          "rulesOperator": "OR",
+          "disableTemplating": false,
+          "fallbackTo404": false,
+          "default": true,
+          "crudKey": "id",
+          "callbacks": []
+        }
+      ],
+      "responseMode": null,
+      "streamingMode": null,
+      "streamingInterval": 0
+    }
+  ],
+  "rootChildren": [
+    {
+      "type": "route",
+      "uuid": "16d88cdc-abcf-4b0d-a4d6-18a3205a9a6c"
+    },
+    {
+      "type": "route",
+      "uuid": "66f51ad5-2ba3-4cb0-b633-c362a9f75836"
+    }
+  ],
+  "proxyMode": false,
+  "proxyHost": "",
+  "proxyRemovePrefix": false,
+  "tlsOptions": {
+    "enabled": false,
+    "type": "CERT",
+    "pfxPath": "",
+    "certPath": "",
+    "keyPath": "",
+    "caPath": "",
+    "passphrase": ""
+  },
+  "cors": true,
+  "headers": [],
+  "proxyReqHeaders": [
+    {
+      "key": "",
+      "value": ""
+    }
+  ],
+  "proxyResHeaders": [
+    {
+      "key": "",
+      "value": ""
+    }
+  ],
+  "data": [],
+  "callbacks": []
+}
diff --git a/.github/workflows/python_build.yml b/.github/workflows/python_build.yml
@@ -84,7 +84,7 @@ jobs:
           uv pip uninstall pandas
           uv run --no-sync pytest -m "not pandas and not integration and not benchmark"
           uv pip install pandas
-  
+
   test-lakefs:
     name: Python Build (Python 3.10 LakeFS Integration tests)
     runs-on: ubuntu-latest
@@ -110,6 +110,64 @@ jobs:
       - name: Run tests
         run: uv run --no-sync pytest -m '(lakefs and integration)' --doctest-modules
 
+  test-unitycatalog-databricks:
+    name: Python Build (Python 3.10 Unity Catalog Integration tests)
+    runs-on: ubuntu-latest
+    env:
+      RUSTFLAGS: "-C debuginfo=1"
+      CARGO_INCREMENTAL: 0
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Setup Environment
+        uses: ./.github/actions/setup-env
+
+      - name: Run Mockoon CLI
+        uses: mockoon/cli-action@v2
+        with:
+          version: "latest"
+          data-file: "./mockoon_data_files/unitycatalog_databricks.json"
+          port: 8080
+
+      - name: Build and install deltalake
+        run: make develop
+
+      - name: Download Data Acceptance Tests (DAT) files
+        run: make setup-dat
+
+      - name: Run tests
+        run: uv run --no-sync pytest -m '(unitycatalog_databricks and integration)' --doctest-modules
+
+  test-unitycatalog-oss:
+    name: Python Build (Python 3.10 Unity Catalog Integration tests)
+    runs-on: ubuntu-latest
+    env:
+      RUSTFLAGS: "-C debuginfo=1"
+      CARGO_INCREMENTAL: 0
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Setup Environment
+        uses: ./.github/actions/setup-env
+
+      - name: Run Mockoon CLI
+        uses: mockoon/cli-action@v2
+        with:
+          version: "latest"
+          data-file: "./mockoon_data_files/unitycatalog_oss.json"
+          port: 8080
+
+      - name: Build and install deltalake
+        run: make develop
+
+      - name: Download Data Acceptance Tests (DAT) files
+        run: make setup-dat
+
+      - name: Run tests
+        run: uv run --no-sync pytest -m '(unitycatalog_oss and integration)' --doctest-modules
+
   test-pyspark:
     name: PySpark Integration Tests
     runs-on: ubuntu-latest
diff --git a/crates/catalog-unity/src/lib.rs b/crates/catalog-unity/src/lib.rs
@@ -473,7 +473,7 @@ impl UnityCatalogBuilder {
             .await
         {
             Ok(s) => s,
-            Err(_e) => panic!("Unable to find the table's storage location."),
+            Err(err) => panic!("Unable to find the table's storage location. {}", err),
         };
         let temp_creds_res = unity_catalog
             .get_temp_table_credentials(catalog_id, database_name, table_name)
diff --git a/python/pyproject.toml b/python/pyproject.toml
@@ -97,6 +97,8 @@ markers = [
     "pandas: marks tests that require pandas",
     "polars: marks tests that require polars",
     "lakefs: marks tests that require lakefs",
+    "unitycatalog_databricks: marks tests that require unitycatalog_databricks",
+    "unitycatalog_oss: marks tests that require unitycatalog_oss",
     "pyspark: marks tests that require pyspark",
 ]
 
diff --git a/python/tests/test_unity_catalog.py b/python/tests/test_unity_catalog.py
@@ -0,0 +1,52 @@
+import os
+
+import pytest
+
+from deltalake import DeltaTable
+
+
+@pytest.mark.unitycatalog_databricks
+@pytest.mark.integration
+@pytest.mark.timeout(timeout=10, method="thread")
+def test_uc_read_deltatable():
+    """Test delta table reads using Unity Catalog URL (uc://)"""
+
+    os.environ["DATABRICKS_WORKSPACE_URL"] = "http://localhost:8080"
+    os.environ["DATABRICKS_ACCESS_TOKEN"] = "123456"
+    os.environ["UNITY_ALLOW_HTTP_URL"] = "true"
+
+    dt = DeltaTable("uc://unity.default.testtable")
+
+    assert dt.is_deltatable(dt.table_uri), True
+    expected = {
+        "value": ["1", "2", "3", "6", "7", "5", "4"],
+        "year": ["2020", "2020", "2020", "2021", "2021", "2021", "2021"],
+        "month": ["1", "2", "2", "12", "12", "12", "4"],
+        "day": ["1", "3", "5", "20", "20", "4", "5"],
+    }
+    assert dt.to_pyarrow_dataset().to_table().to_pydict() == expected
+
+
+@pytest.mark.unitycatalog_oss
+@pytest.mark.integration
+@pytest.mark.timeout(timeout=10, method="thread")
+def test_uc_read_deltatable_failing():
+    """Test delta table reads using Unity Catalog URL (uc://)"""
+
+    os.environ["DATABRICKS_WORKSPACE_URL"] = "http://localhost:8080"
+    os.environ["DATABRICKS_ACCESS_TOKEN"] = "123456"
+    os.environ["UNITY_ALLOW_HTTP_URL"] = "true"
+
+    # @TODO: Currently, this will fail when used with Unity Catalog OSS
+    # mock APIs. Need to add support for slightly different response payloads
+    # of Unity Catalog OSS.
+    with pytest.raises(BaseException):
+        dt = DeltaTable("uc://unity.default.testtable")
+        assert dt.is_deltatable(dt.table_uri), True
+        expected = {
+            "value": ["1", "2", "3", "6", "7", "5", "4"],
+            "year": ["2020", "2020", "2020", "2021", "2021", "2021", "2021"],
+            "month": ["1", "2", "2", "12", "12", "12", "4"],
+            "day": ["1", "3", "5", "20", "20", "4", "5"],
+        }
+        assert dt.to_pyarrow_dataset().to_table().to_pydict() == expected

Original file line number	Diff line number	Diff line change
`@@ -473,7 +473,7 @@ impl UnityCatalogBuilder {`
`473`	`473`	`.await`
`474`	`474`	`{`
`475`	`475`	`Ok(s) => s,`
`476`		`- Err(_e) => panic!("Unable to find the table's storage location."),`
	`476`	`+ Err(err) => panic!("Unable to find the table's storage location. {}", err),`
`477`	`477`	`};`
`478`	`478`	`let temp_creds_res = unity_catalog`
`479`	`479`	`.get_temp_table_credentials(catalog_id, database_name, table_name)`
Original file line number	Diff line number	Diff line change
`@@ -97,6 +97,8 @@ markers = [`
`97`	`97`	`"pandas: marks tests that require pandas",`
`98`	`98`	`"polars: marks tests that require polars",`
`99`	`99`	`"lakefs: marks tests that require lakefs",`
	`100`	`+ "unitycatalog_databricks: marks tests that require unitycatalog_databricks",`
	`101`	`+ "unitycatalog_oss: marks tests that require unitycatalog_oss",`
`100`	`102`	`"pyspark: marks tests that require pyspark",`
`101`	`103`	`]`
`102`	`104`