Add method to upload DuckDB files to Unity Catalog Volume with tests

src/databricks/labs/lakebridge/assessments/dashboards/dashboard_manager.py

@@ -1,3 +1,4 @@
+import io
 import os
 import json
 
@@ -59,3 +60,35 @@ def create_profiler_summary_dashboard(self, extract_file: str | None, source_tec
             run_as_role=None,
             tags=None,
         )
+
+    def upload_duckdb_to_uc_volume(self, local_file_path, volume_path):
+        """
+        Upload a DuckDB file to Unity Catalog Volume
+
+        Args:
+            local_file_path (str): Local path to the DuckDB file
+            volume_path (str): Target path in UC Volume (e.g., '/Volumes/catalog/schema/volume/myfile.duckdb')
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+
+        # Validate inputs
+        if not os.path.exists(local_file_path):
+            logger.error(f"Local file not found: {local_file_path}")
+            return False
+
+        if not volume_path.startswith('/Volumes/'):
+            logger.error("Volume path must start with '/Volumes/'")
+            return False
+
+        try:
+            with open(local_file_path, 'rb') as f:
+                file_bytes = f.read()
+                binary_data = io.BytesIO(file_bytes)
+                self._ws.files.upload(volume_path, binary_data, overwrite = True)
+            logger.info(f"Successfully uploaded {local_file_path} to {volume_path}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to upload file: {str(e)}")
+            return False

tests/integration/assessments/test_dashboard_manager.py

@@ -1,24 +1,56 @@
+import os
+import io
 import pytest
-
-from .utils.profiler_extract_utils import build_mock_synapse_extract
-
-
-@pytest.fixture(scope="module")
-def mock_synapse_profiler_extract():
-    synapse_extract_path = build_mock_synapse_extract("mock_profiler_extract")
-    return synapse_extract_path
-
-
-# Step One:
-# Fetch environment variables for Databricks workspace URL, token, catalog, schema, volume name
-# This will be moved into CLI prompts
-
-# Step Two:
-# Test that the DuckDB file can be uploaded to a target UC Volume
-# TODO: Create class/function for uploading Duck DB file
-
-# Step Three:
-# Test that the job can be deployed to Databricks workspace
-
-# Step Four:
-# Test that the dashboard can be deployed to the workspace
+from unittest.mock import MagicMock, patch
+from databricks.labs.lakebridge.assessments.dashboards.dashboard_manager import DashboardManager
+
+@pytest.fixture
+def mock_workspace_client():
+    return MagicMock()
+
+@pytest.fixture
+def mock_user():
+    return MagicMock()
+
+@pytest.fixture
+def dashboard_manager(mock_workspace_client, mock_user):
+    return DashboardManager(ws=mock_workspace_client, current_user=mock_user)
+
+@patch("os.path.exists")
+def test_upload_duckdb_to_uc_volume_file_not_found(mock_exists, dashboard_manager):
+    mock_exists.return_value = False
+    result = dashboard_manager.upload_duckdb_to_uc_volume("non_existent_file.duckdb", "/Volumes/catalog/schema/volume/myfile.duckdb")
+    assert result is False
+    dashboard_manager._ws.files.upload.assert_not_called()
+
+def test_upload_duckdb_to_uc_volume_invalid_volume_path(dashboard_manager):
+    result = dashboard_manager.upload_duckdb_to_uc_volume("file.duckdb", "invalid_path/myfile.duckdb")
+    assert result is False
+    dashboard_manager._ws.files.upload.assert_not_called()
+
+@patch("os.path.exists")
+@patch("builtins.open", new_callable=MagicMock)
+def test_upload_duckdb_to_uc_volume_success(mock_open, mock_exists, dashboard_manager):
+    mock_exists.return_value = True
+    mock_open.return_value.__enter__.return_value.read.return_value = b"test_data"
+    dashboard_manager._ws.files.upload = MagicMock()
+
+    result = dashboard_manager.upload_duckdb_to_uc_volume("file.duckdb", "/Volumes/catalog/schema/volume/myfile.duckdb")
+    assert result is True
+    dashboard_manager._ws.files.upload.assert_called_once()
+    args, kwargs = dashboard_manager._ws.files.upload.call_args
+    assert args[0] == "/Volumes/catalog/schema/volume/myfile.duckdb"
+    assert isinstance(args[1], io.BytesIO)
+    assert args[1].getvalue() == b"test_data"
+    assert kwargs["overwrite"] is True
+
+@patch("os.path.exists")
+@patch("builtins.open", new_callable=MagicMock)
+def test_upload_duckdb_to_uc_volume_failure(mock_open, mock_exists, dashboard_manager):
+    mock_exists.return_value = True
+    mock_open.return_value.__enter__.return_value.read.return_value = b"test_data"
+    dashboard_manager._ws.files.upload = MagicMock(side_effect=Exception("Upload failed"))
+
+    result = dashboard_manager.upload_duckdb_to_uc_volume("file.duckdb", "/Volumes/catalog/schema/volume/myfile.duckdb")
+    assert result is False
+    dashboard_manager._ws.files.upload.assert_called_once()