def test_missing_table():
error_message = (
"`table` argument cannot be empty. Please provide"
" the name of the table to load or save data to."
)
with pytest.raises(DataSetError, match=error_message):
SparkJDBCDataSet(url="dummy_url", table=None)
def test_missing_url():
error_message = (
"`url` argument cannot be empty. Please provide a JDBC"
" URL of the form ``jdbc:subprotocol:subname``."
)
with pytest.raises(DataSetError, match=error_message):
SparkJDBCDataSet(url=None, table="dummy_table")
def mock_save(arg_dict):
mock_data = mock.Mock()
data_set = SparkJDBCDataSet(**arg_dict)
data_set.save(mock_data)
return mock_data
def mock_load(mock_get_or_create, arg_dict):
spark = mock_get_or_create.return_value
data_set = SparkJDBCDataSet(**arg_dict)
data_set.load()
return spark
