def test_execute_sql_no_results(): global executed_once executed_once = False spark_events = MagicMock() sqlquery = SQLQuery("SHOW TABLES", "take", maxrows=-1, spark_events=spark_events) sqlquery.to_command = MagicMock() sqlquery.to_only_columns_query = MagicMock() result1 = "" result_data = pd.DataFrame([]) session = MagicMock() sqlquery.to_command.return_value.execute.return_value = (True, result1) session.kind = "spark" result = sqlquery.execute(session) assert_frame_equal(result, result_data) sqlquery.to_command.return_value.execute.assert_called_once_with(session) spark_events.emit_sql_execution_start_event.assert_called_once_with( session.guid, session.kind, session.id, sqlquery.guid, sqlquery.samplemethod, sqlquery.maxrows, sqlquery.samplefraction, ) spark_events.emit_sql_execution_end_event.assert_called_once_with( session.guid, session.kind, session.id, sqlquery.guid, sqlquery.to_command.return_value.guid, True, "", "" )
def test_execute_sql(): spark_events = MagicMock() sqlquery = SQLQuery("HERE IS THE QUERY", "take", 100, 0.2, spark_events=spark_events) sqlquery.to_command = MagicMock(return_value=MagicMock()) result = """{"z":100, "nullv":null, "y":50} {"z":25, "nullv":null, "y":10}""" sqlquery.to_command.return_value.execute = MagicMock(return_value=(True, result)) result_data = pd.DataFrame([{ 'z': 100, "nullv": None, 'y': 50 }, { 'z': 25, "nullv": None, 'y': 10 }], columns=['z', "nullv", 'y']) session = MagicMock() session.kind = "pyspark" result = sqlquery.execute(session) assert_frame_equal(result, result_data) sqlquery.to_command.return_value.execute.assert_called_once_with(session) spark_events.emit_sql_execution_start_event.assert_called_once_with( session.guid, session.kind, session.id, sqlquery.guid, 'take', 100, 0.2) spark_events.emit_sql_execution_end_event.assert_called_once_with( session.guid, session.kind, session.id, sqlquery.guid, sqlquery.to_command.return_value.guid, True, '', '')
def test_execute_sql_no_results(): global executed_once executed_once = False spark_events = MagicMock() sqlquery = SQLQuery("SHOW TABLES", "take", maxrows=-1, spark_events=spark_events) sqlquery.to_command = MagicMock() sqlquery.to_only_columns_query = MagicMock() result1 = "" result_data = pd.DataFrame([]) session = MagicMock() sqlquery.to_command.return_value.execute.return_value = ( True, result1, MIMETYPE_TEXT_PLAIN) session.kind = "spark" result = sqlquery.execute(session) assert_frame_equal(result, result_data) sqlquery.to_command.return_value.execute.assert_called_once_with(session) spark_events.emit_sql_execution_start_event.assert_called_once_with( session.guid, session.kind, session.id, sqlquery.guid, sqlquery.samplemethod, sqlquery.maxrows, sqlquery.samplefraction) spark_events.emit_sql_execution_end_event.assert_called_once_with( session.guid, session.kind, session.id, sqlquery.guid, sqlquery.to_command.return_value.guid, True, "", "")
def test_execute_sql_failure_emits_event(): spark_events = MagicMock() sqlquery = SQLQuery("HERE IS THE QUERY", "take", 100, 0.2, spark_events) sqlquery.to_command = MagicMock() sqlquery.to_command.return_value.execute = MagicMock(side_effect=ValueError('yo')) session = MagicMock() session.kind = "pyspark" try: result = sqlquery.execute(session) assert False except ValueError: sqlquery.to_command.return_value.execute.assert_called_once_with(session) spark_events.emit_sql_execution_end_event.assert_called_once_with(session.guid, session.kind, session.id, sqlquery.guid, sqlquery.to_command.return_value.guid, False, 'ValueError', 'yo')
def test_execute_sql(): spark_events = MagicMock() sqlquery = SQLQuery("HERE IS THE QUERY", "take", 100, 0.2, spark_events=spark_events) sqlquery.to_command = MagicMock(return_value=MagicMock()) result = """{"z":100, "nullv":null, "y":50} {"z":25, "nullv":null, "y":10}""" sqlquery.to_command.return_value.execute = MagicMock(return_value=(True, result)) result_data = pd.DataFrame([{'z': 100, "nullv": None, 'y': 50}, {'z':25, "nullv":None, 'y':10}], columns=['z', "nullv", 'y']) session = MagicMock() session.kind = "pyspark" result = sqlquery.execute(session) assert_frame_equal(result, result_data) sqlquery.to_command.return_value.execute.assert_called_once_with(session) spark_events.emit_sql_execution_start_event._assert_called_once_with(session.guid, session.kind, session.id, sqlquery.guid) spark_events.emit_sql_execution_end_event._assert_called_once_with(session.guid, session.kind, session.id, sqlquery.guid, sqlquery.to_command.return_value.guid, True, "", "")
def test_to_command_pyspark3(): variable_name = "var_name" sqlquery = SQLQuery("Query") sqlquery._pyspark_command = MagicMock(return_value=MagicMock()) sqlquery.to_command("pyspark3", variable_name) sqlquery._pyspark_command.assert_called_with(variable_name, False)