def test_execute_code_pandas_pyspark_livy(self):
mock_spark_session = MagicMock()
client = PandasScalaLivyClient(mock_spark_session, 10)
command = "command"
client.execute(command)
mock_spark_session.wait_for_state.assert_called_with("idle")
mock_spark_session.execute.assert_called_with(command)
