mocking unit tests

Author: ahosler

tests/operators/common/test_load_data.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 from typing import Union
 
-# Copyright (c) 2024 Oracle and/or its affiliates.
+# Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 import pytest
 from ads.opctl.operator.lowcode.common.utils import (
@@ -13,23 +13,20 @@
 import pandas as pd
 
 mock_secret = {
-    'user_name': 'mock_user',
-    'password': 'mock_password',
-    'service_name': 'mock_service_name'
+    "user_name": "mock_user",
+    "password": "mock_password",
+    "service_name": "mock_service_name",
 }
 
 mock_connect_args = {
-    'user': 'mock_user',
-    'password': 'mock_password',
-    'service_name': 'mock_service_name',
-    'dsn': 'mock_dsn'
+    "user": "mock_user",
+    "password": "mock_password",
+    "service_name": "mock_service_name",
+    "dsn": "mock_dsn",
 }
 
 # Mock data for testing
-mock_data = pd.DataFrame({
-    'id': [1, 2, 3],
-    'name': ['Alice', 'Bob', 'Charlie']
-})
+mock_data = pd.DataFrame({"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"]})
 
 mock_db_connection = MagicMock()
 
@@ -42,7 +39,9 @@ def mock_oracledb_connect_failure(*args, **kwargs):
 
 
 def mock_oracledb_connect(**kwargs):
-    assert kwargs == mock_connect_args, f"Expected connect_args {mock_connect_args}, but got {kwargs}"
+    assert (
+        kwargs == mock_connect_args
+    ), f"Expected connect_args {mock_connect_args}, but got {kwargs}"
     return mock_db_connection
 
 
@@ -67,37 +66,50 @@ def load_secret_fail(*args, **kwargs):
 class TestDataLoad(unittest.TestCase):
     def setUp(self):
         self.data_spec = Mock(spec=InputData)
-        self.data_spec.connect_args = {
-            'dsn': 'mock_dsn'
-        }
-        self.data_spec.vault_secret_id = 'mock_secret_id'
-        self.data_spec.table_name = 'mock_table_name'
+        self.data_spec.connect_args = {"dsn": "mock_dsn"}
+        self.data_spec.vault_secret_id = "mock_secret_id"
+        self.data_spec.table_name = "mock_table_name"
         self.data_spec.url = None
         self.data_spec.format = None
         self.data_spec.columns = None
         self.data_spec.limit = None
+        self.data_spec.data = None
 
     def testLoadSecretAndDBConnection(self):
-        with patch('ads.secrets.ADBSecretKeeper.load_secret', side_effect=MockADBSecretKeeper.load_secret):
-            with patch('oracledb.connect', side_effect=mock_oracledb_connect):
-                with patch('pandas.read_sql', return_value=mock_data) as mock_read_sql:
+        with patch(
+            "ads.secrets.ADBSecretKeeper.load_secret",
+            side_effect=MockADBSecretKeeper.load_secret,
+        ):
+            with patch("oracledb.connect", side_effect=mock_oracledb_connect):
+                with patch("pandas.read_sql", return_value=mock_data) as mock_read_sql:
                     data = load_data(self.data_spec)
-                    mock_read_sql.assert_called_once_with(f"SELECT * FROM {self.data_spec.table_name}",
-                                                          mock_db_connection)
+                    mock_read_sql.assert_called_once_with(
+                        f"SELECT * FROM {self.data_spec.table_name}", mock_db_connection
+                    )
                     pd.testing.assert_frame_equal(data, mock_data)
 
     def testLoadVaultFailure(self):
-        with patch('ads.secrets.ADBSecretKeeper.load_secret', side_effect=MockADBSecretKeeper.load_secret_fail):
+        with patch(
+            "ads.secrets.ADBSecretKeeper.load_secret",
+            side_effect=MockADBSecretKeeper.load_secret_fail,
+        ):
             with pytest.raises(Exception) as e:
                 load_data(self.data_spec)
 
         expected_msg = f"Could not retrieve database credentials from vault {self.data_spec.vault_secret_id}: {load_secret_err_msg}"
-        assert str(e.value) == expected_msg, f"Expected exception message '{expected_msg}', but got '{str(e)}'"
+        assert (
+            str(e.value) == expected_msg
+        ), f"Expected exception message '{expected_msg}', but got '{str(e)}'"
 
     def testDBConnectionFailure(self):
-        with patch('ads.secrets.ADBSecretKeeper.load_secret', side_effect=MockADBSecretKeeper.load_secret):
-            with patch('oracledb.connect', side_effect=mock_oracledb_connect_failure):
+        with patch(
+            "ads.secrets.ADBSecretKeeper.load_secret",
+            side_effect=MockADBSecretKeeper.load_secret,
+        ):
+            with patch("oracledb.connect", side_effect=mock_oracledb_connect_failure):
                 with pytest.raises(Exception) as e:
                     load_data(self.data_spec)
 
-        assert str(e.value) == db_connect_err_msg , f"Expected exception message '{db_connect_err_msg }', but got '{str(e)}'"
+        assert (
+            str(e.value) == db_connect_err_msg
+        ), f"Expected exception message '{db_connect_err_msg }', but got '{str(e)}'"