Skip to content

Commit

Permalink
test case
Browse files Browse the repository at this point in the history
  • Loading branch information
v-chen_data committed Sep 25, 2024
1 parent 4c55fa6 commit d3eeb88
Showing 1 changed file with 35 additions and 0 deletions.
35 changes: 35 additions & 0 deletions tests/a_scripts/data_prep/test_convert_delta_to_json.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
# Copyright 2022 MosaicML LLM Foundry authors
# SPDX-License-Identifier: Apache-2.0

import sys
import unittest
from argparse import Namespace
from typing import Any
from unittest.mock import MagicMock, mock_open, patch

from llmfoundry.command_utils.data_prep.convert_delta_to_json import (
InsufficientPermissionsError,
download,
fetch_DT,
format_tablename,
Expand All @@ -17,6 +19,39 @@

class TestConvertDeltaToJsonl(unittest.TestCase):

def test_run_query_dbconnect_insufficient_permissions(self):
error_message = (
'[INSUFFICIENT_PERMISSIONS] Insufficient privileges: User does not have USE SCHEMA '
"on Schema 'main.oogabooga'. SQLSTATE: 42501"
)

class MockAnalysisException(Exception):

def __init__(self, message: str):
self.message = message

with patch.dict('sys.modules', {'pyspark.errors': MagicMock()}):
sys.modules[
'pyspark.errors'
].AnalysisException = MockAnalysisException # pyright: ignore

mock_spark = MagicMock()
mock_spark.sql.side_effect = MockAnalysisException(error_message)

with self.assertRaises(InsufficientPermissionsError) as context:
run_query(
'SELECT * FROM table',
method='dbconnect',
cursor=None,
spark=mock_spark,
)

self.assertIn(
'using the schema main.oogabooga',
str(context.exception),
)
mock_spark.sql.assert_called_once_with('SELECT * FROM table')

@patch(
'databricks.sql.connect',
)
Expand Down

0 comments on commit d3eeb88

Please sign in to comment.