def test_get_dbutils_from_ipython(self, mocker):
ipython_mock = mocker.Mock()
ipython_mock.get_ipython.return_value.user_ns = {
"dbutils": "dbutils_from_ipython"
}
mocker.patch.dict("sys.modules", {"IPython": ipython_mock})
assert _get_dbutils("spark") == "dbutils_from_ipython"
ipython_mock.get_ipython.assert_called_once_with()
def test_get_dbutils_no_modules(self, mocker):
mocker.patch("kedro.extras.datasets.spark.spark_dataset.globals",
return_value={})
mocker.patch.dict("sys.modules", {"pyspark": None, "IPython": None})
assert _get_dbutils("spark") is None
def test_get_dbutils_from_pyspark(self, mocker):
dbutils_mock = mocker.Mock()
dbutils_mock.DBUtils.return_value = "dbutils_from_pyspark"
mocker.patch.dict("sys.modules", {"pyspark.dbutils": dbutils_mock})
assert _get_dbutils("spark") == "dbutils_from_pyspark"
dbutils_mock.DBUtils.assert_called_once_with("spark")
def test_get_dbutils_from_globals(self, mocker):
mocker.patch(
"kedro.extras.datasets.spark.spark_dataset.globals",
return_value={"dbutils": "dbutils_from_globals"},
)
assert _get_dbutils("spark") == "dbutils_from_globals"
