def test_execute_select_query(mocked_client, mocker) -> None:
class Users(declarative_base()):
__tablename__ = 'users'
id = Column(Integer, primary_key=True, autoincrement=True)
name = Column(String(255, collation='utf8_unicode_ci'), default=None)
mocked_client.execute_statement.return_value = {
'numberOfRecordsUpdated': 0,
'records': [[{
'longValue': 1
}, {
'stringValue': 'cat'
}]],
}
data_api = DataAPI(
resource_arn='arn:aws:rds:dummy',
secret_arn='dummy',
database='test',
client=mocked_client,
)
result = data_api.execute(Query(Users).filter(Users.id == 1))
assert len(result) == 1
assert list(result[0]) == [1, 'cat']
assert mocked_client.execute_statement.call_args == mocker.call(
continueAfterTimeout=True,
database='test',
includeResultMetadata=True,
resourceArn='arn:aws:rds:dummy',
secretArn='dummy',
sql='SELECT users.id, users.name \nFROM users \nWHERE users.id = 1',
)
def test_execute_insert(mocked_client, mocker) -> None:
mocked_client.execute_statement.return_value = {
'generatedFields': [],
'numberOfRecordsUpdated': 1,
}
data_api = DataAPI(
resource_arn='arn:aws:rds:dummy',
secret_arn='dummy',
database='test',
client=mocked_client,
)
results = data_api.execute(
"insert into pets values(1, 'cat')", transaction_id='abc'
)
assert results.generated_fields == []
assert results.number_of_records_updated == 1
assert mocked_client.execute_statement.call_args == mocker.call(
continueAfterTimeout=True,
includeResultMetadata=True,
resourceArn='arn:aws:rds:dummy',
secretArn='dummy',
transactionId='abc',
sql="insert into pets values(1, 'cat')",
database='test',
)
def test_execute_query_process_result_value(mocked_client, mocker) -> None:
mocked_client.execute_statement.return_value = {
'numberOfRecordsUpdated':
0,
'records': [[{
'longValue': 1
}, {
'stringValue': 'cat'
}]],
'columnMetadata': [
{
"arrayBaseColumnType": 0,
"isAutoIncrement": False,
"isCaseSensitive": False,
"isCurrency": False,
"isSigned": True,
"label": "id",
"name": "id",
"nullable": 1,
"precision": 11,
"scale": 0,
"schemaName": "",
"tableName": "pets",
"type": 4,
"typeName": "INT",
},
{
"arrayBaseColumnType": 0,
"isAutoIncrement": False,
"isCaseSensitive": False,
"isCurrency": False,
"isSigned": False,
"label": "name",
"name": "name",
"nullable": 1,
"precision": 255,
"scale": 0,
"schemaName": "",
"tableName": "pets",
"type": 12,
"typeName": "VARCHAR",
},
],
}
data_api = DataAPI(
resource_arn='arn:aws:rds:dummy',
secret_arn='dummy',
database='test',
client=mocked_client,
)
assert list(data_api.execute(Query(Pets))[0]) == [1, 'my_type_cat']
assert mocked_client.execute_statement.call_args == mocker.call(
continueAfterTimeout=True,
database='test',
includeResultMetadata=True,
resourceArn='arn:aws:rds:dummy',
secretArn='dummy',
sql="""SELECT pets.id, pets.name
FROM pets""",
)
def test_execute_insert_parameters(mocked_client, mocker) -> None:
mocked_client.execute_statement.return_value = {
'generatedFields': [],
'numberOfRecordsUpdated': 1,
}
data_api = DataAPI(
resource_arn='dummy', secret_arn='dummy', database='test', client=mocked_client
)
results = data_api.execute(
"insert into pets values(:id, :name)",
{'id': 1, 'name': 'cat'},
transaction_id='abc',
)
assert results.generated_fields == []
assert results.number_of_records_updated == 1
assert mocked_client.execute_statement.call_args == mocker.call(
continueAfterTimeout=True,
includeResultMetadata=True,
parameters=[
{'name': 'id', 'value': {'longValue': 1}},
{'name': 'name', 'value': {'stringValue': 'cat'}},
],
resourceArn='dummy',
secretArn='dummy',
transactionId='abc',
sql="insert into pets values(:id, :name)",
database='test',
)
def test_execute_select_as_model(mocked_client, mocker) -> None:
mocked_client.execute_statement.return_value = {
"columnMetadata": [
{
"arrayBaseColumnType": 0,
"isAutoIncrement": False,
"isCaseSensitive": False,
"isCurrency": False,
"isSigned": True,
"label": "id",
"name": "id",
"nullable": 1,
"precision": 11,
"scale": 0,
"schemaName": "",
"tableName": "users",
"type": 4,
"typeName": "INT",
},
{
"arrayBaseColumnType": 0,
"isAutoIncrement": False,
"isCaseSensitive": False,
"isCurrency": False,
"isSigned": False,
"label": "name",
"name": "name",
"nullable": 1,
"precision": 255,
"scale": 0,
"schemaName": "",
"tableName": "users",
"type": 12,
"typeName": "VARCHAR",
},
],
'numberOfRecordsUpdated': 0,
'records': [[{'longValue': 1}, {'stringValue': 'cat'}]],
}
data_api = DataAPI(
resource_arn='dummy', secret_arn='dummy', database='test', client=mocked_client
)
class Pet(BaseModel):
id: int
name: str
result = data_api.execute("select * from pets")
assert len(result) == 1
assert result[0].model(Pet) == Pet(name='cat', id=1)
assert mocked_client.execute_statement.call_args == mocker.call(
continueAfterTimeout=True,
database='test',
includeResultMetadata=True,
resourceArn='dummy',
secretArn='dummy',
sql='select * from pets',
)
def test_rollback(mocked_client, mocker) -> None:
mocked_client.rollback_transaction.return_value = {'transactionStatus': 'abc'}
data_api = DataAPI(
resource_arn='dummy', secret_arn='dummy', database='test', client=mocked_client
)
assert data_api.rollback(transaction_id='abc') == 'abc'
assert mocked_client.rollback_transaction.call_args == mocker.call(
resourceArn='dummy', transactionId='abc', secretArn='dummy'
)
def test_transaction_status(mocked_client) -> None:
data_api = DataAPI(
resource_arn='arn:aws:rds:dummy',
secret_arn='dummy',
transaction_id='abc',
client=mocked_client,
)
data_api._transaction_status = 'dummy status'
assert data_api.transaction_status == 'dummy status'
def test_execute_insert_parameter_set_invalid(mocked_client, mocker) -> None:
data_api = DataAPI(
resource_arn='dummy', secret_arn='dummy', database='test', client=mocked_client
)
with pytest.raises(ValidationError):
data_api.batch_execute(
"insert into test.pets values (:id , :name)", {'id': 3, 'name': 'bird'}
)
def test_begin(mocked_client, mocker) -> None:
mocked_client.begin_transaction.return_value = {'transactionId': 'abc'}
data_api = DataAPI(
resource_arn='dummy', secret_arn='dummy', database='test', client=mocked_client
)
assert data_api.begin(schema='schema') == 'abc'
assert mocked_client.begin_transaction.call_args == mocker.call(
database='test', resourceArn='dummy', schema='schema', secretArn='dummy'
)
def test_execute_insert_parameter_set_invalid_1(mocked_client, mocker) -> None:
mocked_client.batch_execute_statement.side_effect = Exception('Invalid Request')
data_api = DataAPI(
resource_arn='arn:aws:rds:dummy',
secret_arn='dummy',
database='test',
client=mocked_client,
transaction_id='12345',
)
with pytest.raises(Exception):
data_api.batch_execute(
"insert into test.pets values (:id , :name)",
[{'id': 3, 'invalid': 'bird'}],
)
def test_not_found_resource_arn_and_resource_arn(mocker, mocked_client) -> None:
mock_client = mocker.Mock()
mock_client.describe_db_clusters.return_value = {
'DBClusters': [{'DBClusterArn': 'arn:aws:rds:dummy'}]
}
with pytest.raises(DataAPIError, match='Not Found resource_arn.'):
DataAPI(secret_arn='dummy', client=mock_client, rds_client=mock_client)
def test_execute_select(mocked_client, mocker) -> None:
mocked_client.execute_statement.return_value = {
'numberOfRecordsUpdated': 0,
'records': [[{'longValue': 1}, {'stringValue': 'cat'}]],
}
data_api = DataAPI(
resource_arn='dummy', secret_arn='dummy', database='test', client=mocked_client
)
assert list(data_api.execute("select * from pets")[0]) == [1, 'cat']
assert mocked_client.execute_statement.call_args == mocker.call(
continueAfterTimeout=True,
database='test',
includeResultMetadata=True,
resourceArn='dummy',
secretArn='dummy',
sql='select * from pets',
)
def test_transaction(mocked_client) -> None:
data_api = DataAPI(
resource_arn='dummy',
secret_arn='dummy',
transaction_id='abc',
client=mocked_client,
)
assert data_api.transaction_id == 'abc'
def test_resource_arn(mocker, mocked_client) -> None:
mock_client = mocker.Mock()
mock_client.describe_db_clusters.return_value = {
'DBClusters': [{'DBClusterArn': 'arn:aws:rds:dummy'}]
}
data_api: DataAPI = DataAPI(
resource_name='dummy',
secret_arn='dummy',
client=mock_client,
rds_client=mock_client,
)
assert data_api.resource_arn == 'arn:aws:rds:dummy'
mocked_client.return_value = mock_client
data_api: DataAPI = DataAPI(
resource_name='dummy', secret_arn='dummy', client=mock_client
)
assert data_api.resource_arn == 'arn:aws:rds:dummy'
def test_execute_insert_parameter_set(mocked_client, mocker) -> None:
mocked_client.batch_execute_statement.return_value = {
'updateResults': [
{'generatedFields': [{'longValue': 3}]},
{'generatedFields': [{'longValue': 4}]},
]
}
mocked_client.begin_transaction.return_value = {'transactionId': '12345'}
data_api = DataAPI(
resource_arn='arn:aws:rds:dummy',
secret_arn='dummy',
database='test',
client=mocked_client,
)
results = data_api.batch_execute(
"insert into test.pets values (:id , :name)",
[{'id': 3, 'name': 'bird'}, {'id': 4, 'name': 'lion'}],
)
assert len(results) == 2
assert results[0].generated_fields == [3]
assert results[0].generated_fields_first == 3
assert results[1].generated_fields == [4]
assert results[1].generated_fields_first == 4
assert mocked_client.batch_execute_statement.call_args == mocker.call(
resourceArn='arn:aws:rds:dummy',
secretArn='dummy',
sql="insert into test.pets values (:id , :name)",
parameterSets=[
[
{'name': 'id', 'value': {'longValue': 3}},
{'name': 'name', 'value': {'stringValue': 'bird'}},
],
[
{'name': 'id', 'value': {'longValue': 4}},
{'name': 'name', 'value': {'stringValue': 'lion'}},
],
],
database='test',
transactionId='12345',
)
def test_with_statement(mocked_client) -> None:
mocked_client.begin_transaction.return_value = {'transactionId': 'abc'}
with DataAPI(
resource_arn='dummy', secret_arn='dummy', database='test', client=mocked_client
):
mocked_client.begin_transaction.assert_called_once_with(
database='test', resourceArn='dummy', secretArn='dummy'
)
mocked_client.commit_transaction.assert_called_once_with(
resourceArn='dummy', secretArn='dummy', transactionId='abc'
)
def test_with_statement_custom_exception(mocked_client, mocker) -> None:
class CustomError(Exception):
pass
mocked_client.begin_transaction.return_value = {'transactionId': 'abc'}
with pytest.raises(CustomError):
with DataAPI(
resource_arn='arn:aws:rds:dummy',
secret_arn='dummy',
database='test',
client=mocked_client,
rollback_exception=CustomError,
):
mocked_client.begin_transaction.assert_called_once_with(
database='test',
resourceArn='arn:aws:rds:dummy',
secretArn='dummy')
raise CustomError('error')
mocked_client.rollback_transaction.assert_called_once_with(
resourceArn='arn:aws:rds:dummy',
secretArn='dummy',
transactionId='abc')
second_mocked_client = mocker.patch('boto3.client')
second_mocked_client.begin_transaction.return_value = {
'transactionId': 'abc'
}
with pytest.raises(Exception):
with DataAPI(
resource_arn='arn:aws:rds:dummy',
secret_arn='dummy',
database='test',
client=second_mocked_client,
rollback_exception=CustomError,
):
second_mocked_client.begin_transaction.assert_called_once_with(
database='test',
resourceArn='arn:aws:rds:dummy',
secretArn='dummy')
raise Exception('error')
second_mocked_client.rollback_transaction.assert_not_called()
def test_found_resource_arn_and_resource_arn(mocker, mocked_client) -> None:
mock_client = mocker.Mock()
mock_client.describe_db_clusters.return_value = {
'DBClusters': [{'DBClusterArn': 'arn:aws:rds:dummy'}]
}
with pytest.raises(
DataAPIError,
match='resource_name should be set without resource_arn. resource_arn: arn:aws:rds:dummy, resource_name: dummy',
):
DataAPI(
resource_arn='arn:aws:rds:dummy',
resource_name='dummy',
secret_arn='dummy',
client=mock_client,
rds_client=mock_client,
)
def test_with_statement_exception(mocked_client) -> None:
mocked_client.begin_transaction.return_value = {'transactionId': 'abc'}
with pytest.raises(Exception):
with DataAPI(
resource_arn='dummy',
secret_arn='dummy',
database='test',
client=mocked_client,
):
mocked_client.begin_transaction.assert_called_once_with(
database='test', resourceArn='dummy', secretArn='dummy'
)
raise Exception('error')
mocked_client.rollback_transaction.assert_called_once_with(
resourceArn='dummy', secretArn='dummy', transactionId='abc'
)
def test_client(mocker) -> None:
mock_client = mocker.Mock()
data_api: DataAPI = DataAPI(resource_arn='arn:aws:rds:dummy',
secret_arn='dummy',
client=mock_client)
assert data_api.client == mock_client
def add_user(data_api: DataAPI, id_, name):
data_api.execute(f"insert into pets values({id_}, {name})")
