1
1
#!/usr/bin/env python
2
2
from typing import Union
3
3
4
- # Copyright (c) 2024 Oracle and/or its affiliates.
4
+ # Copyright (c) 2024, 2025 Oracle and/or its affiliates.
5
5
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
6
6
import pytest
7
7
from ads .opctl .operator .lowcode .common .utils import (
13
13
import pandas as pd
14
14
15
15
mock_secret = {
16
- ' user_name' : ' mock_user' ,
17
- ' password' : ' mock_password' ,
18
- ' service_name' : ' mock_service_name'
16
+ " user_name" : " mock_user" ,
17
+ " password" : " mock_password" ,
18
+ " service_name" : " mock_service_name" ,
19
19
}
20
20
21
21
mock_connect_args = {
22
- ' user' : ' mock_user' ,
23
- ' password' : ' mock_password' ,
24
- ' service_name' : ' mock_service_name' ,
25
- ' dsn' : ' mock_dsn'
22
+ " user" : " mock_user" ,
23
+ " password" : " mock_password" ,
24
+ " service_name" : " mock_service_name" ,
25
+ " dsn" : " mock_dsn" ,
26
26
}
27
27
28
28
# Mock data for testing
29
- mock_data = pd .DataFrame ({
30
- 'id' : [1 , 2 , 3 ],
31
- 'name' : ['Alice' , 'Bob' , 'Charlie' ]
32
- })
29
+ mock_data = pd .DataFrame ({"id" : [1 , 2 , 3 ], "name" : ["Alice" , "Bob" , "Charlie" ]})
33
30
34
31
mock_db_connection = MagicMock ()
35
32
@@ -42,7 +39,9 @@ def mock_oracledb_connect_failure(*args, **kwargs):
42
39
43
40
44
41
def mock_oracledb_connect (** kwargs ):
45
- assert kwargs == mock_connect_args , f"Expected connect_args { mock_connect_args } , but got { kwargs } "
42
+ assert (
43
+ kwargs == mock_connect_args
44
+ ), f"Expected connect_args { mock_connect_args } , but got { kwargs } "
46
45
return mock_db_connection
47
46
48
47
@@ -67,37 +66,50 @@ def load_secret_fail(*args, **kwargs):
67
66
class TestDataLoad (unittest .TestCase ):
68
67
def setUp (self ):
69
68
self .data_spec = Mock (spec = InputData )
70
- self .data_spec .connect_args = {
71
- 'dsn' : 'mock_dsn'
72
- }
73
- self .data_spec .vault_secret_id = 'mock_secret_id'
74
- self .data_spec .table_name = 'mock_table_name'
69
+ self .data_spec .connect_args = {"dsn" : "mock_dsn" }
70
+ self .data_spec .vault_secret_id = "mock_secret_id"
71
+ self .data_spec .table_name = "mock_table_name"
75
72
self .data_spec .url = None
76
73
self .data_spec .format = None
77
74
self .data_spec .columns = None
78
75
self .data_spec .limit = None
76
+ self .data_spec .data = None
79
77
80
78
def testLoadSecretAndDBConnection (self ):
81
- with patch ('ads.secrets.ADBSecretKeeper.load_secret' , side_effect = MockADBSecretKeeper .load_secret ):
82
- with patch ('oracledb.connect' , side_effect = mock_oracledb_connect ):
83
- with patch ('pandas.read_sql' , return_value = mock_data ) as mock_read_sql :
79
+ with patch (
80
+ "ads.secrets.ADBSecretKeeper.load_secret" ,
81
+ side_effect = MockADBSecretKeeper .load_secret ,
82
+ ):
83
+ with patch ("oracledb.connect" , side_effect = mock_oracledb_connect ):
84
+ with patch ("pandas.read_sql" , return_value = mock_data ) as mock_read_sql :
84
85
data = load_data (self .data_spec )
85
- mock_read_sql .assert_called_once_with (f"SELECT * FROM { self .data_spec .table_name } " ,
86
- mock_db_connection )
86
+ mock_read_sql .assert_called_once_with (
87
+ f"SELECT * FROM { self .data_spec .table_name } " , mock_db_connection
88
+ )
87
89
pd .testing .assert_frame_equal (data , mock_data )
88
90
89
91
def testLoadVaultFailure (self ):
90
- with patch ('ads.secrets.ADBSecretKeeper.load_secret' , side_effect = MockADBSecretKeeper .load_secret_fail ):
92
+ with patch (
93
+ "ads.secrets.ADBSecretKeeper.load_secret" ,
94
+ side_effect = MockADBSecretKeeper .load_secret_fail ,
95
+ ):
91
96
with pytest .raises (Exception ) as e :
92
97
load_data (self .data_spec )
93
98
94
99
expected_msg = f"Could not retrieve database credentials from vault { self .data_spec .vault_secret_id } : { load_secret_err_msg } "
95
- assert str (e .value ) == expected_msg , f"Expected exception message '{ expected_msg } ', but got '{ str (e )} '"
100
+ assert (
101
+ str (e .value ) == expected_msg
102
+ ), f"Expected exception message '{ expected_msg } ', but got '{ str (e )} '"
96
103
97
104
def testDBConnectionFailure (self ):
98
- with patch ('ads.secrets.ADBSecretKeeper.load_secret' , side_effect = MockADBSecretKeeper .load_secret ):
99
- with patch ('oracledb.connect' , side_effect = mock_oracledb_connect_failure ):
105
+ with patch (
106
+ "ads.secrets.ADBSecretKeeper.load_secret" ,
107
+ side_effect = MockADBSecretKeeper .load_secret ,
108
+ ):
109
+ with patch ("oracledb.connect" , side_effect = mock_oracledb_connect_failure ):
100
110
with pytest .raises (Exception ) as e :
101
111
load_data (self .data_spec )
102
112
103
- assert str (e .value ) == db_connect_err_msg , f"Expected exception message '{ db_connect_err_msg } ', but got '{ str (e )} '"
113
+ assert (
114
+ str (e .value ) == db_connect_err_msg
115
+ ), f"Expected exception message '{ db_connect_err_msg } ', but got '{ str (e )} '"
0 commit comments