def test_execute_failure_wait_for_session_emits_event():
    spark_events = MagicMock()
    kind = SESSION_KIND_SPARK
    http_client = MagicMock()
    http_client.post_session.return_value = tls.TestLivySession.session_create_json
    http_client.post_statement.return_value = tls.TestLivySession.post_statement_json
    http_client.get_session.return_value = tls.TestLivySession.ready_sessions_json
    http_client.get_statement.return_value = tls.TestLivySession.ready_statement_json
    conf.override_all({
        "status_sleep_seconds": 0.01,
        "statement_sleep_seconds": 0.01
    })
    session = _create_session(kind=kind, http_client=http_client)
    conf.override_all({})
    session.start(create_sql_context=False)
    session.wait_for_idle = MagicMock(side_effect=ValueError("yo"))
    command = Command("command", spark_events=spark_events)

    try:
        result = command.execute(session)
        assert False
    except ValueError as e:
        spark_events.emit_statement_execution_start_event._assert_called_once_with(session.guid, session.kind,
                                                                                   session.id, command.guid)
        spark_events.emit_statement_execution_start_event._assert_called_once_with(session.guid, session.kind,
                                                                                   session.id, command.guid,
                                                                                   -1, False, "ValueError", "yo")
        assert_equals(e, session.wait_for_idle.side_effect)

def test_execute_failure_post_statement_emits_event():
    spark_events = MagicMock()
    kind = SESSION_KIND_SPARK
    http_client = MagicMock()
    http_client.get_statement.return_value = tls.TestLivySession.ready_statement_json
    conf.override_all({"status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01})
    session = _create_session(kind=kind, http_client=http_client)
    session.wait_for_idle = MagicMock()
    conf.override_all({})
    session.start()
    session.wait_for_idle = MagicMock()
    command = Command("command", spark_events=spark_events)

    http_client.post_statement.side_effect = KeyError("Something bad happened here")
    try:
        result = command.execute(session)
        assert False
    except KeyError as e:
        spark_events.emit_statement_execution_start_event.assert_called_once_with(
            session.guid, session.kind, session.id, command.guid
        )
        spark_events.emit_statement_execution_end_event._assert_called_once_with(
            session.guid, session.kind, session.id, command.guid, -1, False, "KeyError", "Something bad happened here"
        )
        assert_equals(e, http_client.post_statement.side_effect)

def test_execute():
    spark_events = MagicMock()
    kind = SESSION_KIND_SPARK
    http_client = MagicMock()
    http_client.post_session.return_value = tls.TestLivySession.session_create_json
    http_client.post_statement.return_value = tls.TestLivySession.post_statement_json
    http_client.get_session.return_value = tls.TestLivySession.ready_sessions_json
    http_client.get_statement.return_value = tls.TestLivySession.ready_statement_json
    conf.override_all({"status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01})
    session = _create_session(kind=kind, http_client=http_client)
    conf.override_all({})
    session.start()
    command = Command("command", spark_events=spark_events)

    result = command.execute(session)

    http_client.post_statement.assert_called_with(0, {"code": command.code})
    http_client.get_statement.assert_called_with(0, 0)
    assert result[0]
    assert_equals(tls.TestLivySession.pi_result, result[1])
    spark_events.emit_statement_execution_start_event.assert_called_once_with(
        session.guid, session.kind, session.id, command.guid
    )
    spark_events.emit_statement_execution_end_event.assert_called_once_with(
        session.guid, session.kind, session.id, command.guid, 0, True, "", ""
    )

    def test_delete_session_emits_start_failed_end_when_delete_throws(self):
        self.http_client.delete_session.side_effect = ValueError
        self.http_client.get_session.return_value = self.ready_sessions_json
        self.http_client.get_statement.return_value = self.ready_statement_json

        conf.override_all({
            "status_sleep_seconds": 0.01,
            "statement_sleep_seconds": 0.01
        })
        session = self._create_session()
        conf.override_all({})
        session.start()
        session.id = 0
        end_id = session.id
        end_status = constants.BUSY_SESSION_STATUS
        session.status = end_status

        try:
            session.delete()
            assert False
        except ValueError:
            pass

        self.spark_events.emit_session_deletion_start_event.assert_called_once_with(
            session.guid, session.kind, end_id, end_status)
        self.spark_events.emit_session_deletion_end_event.assert_called_once_with(
            session.guid, session.kind, end_id, end_status, False, "ValueError", "")

 def test_constructor_throws_status_sleep_seconds(self):
     conf.override_all({
         "status_sleep_seconds": 0,
         "statement_sleep_seconds": 2
     })
     self._create_session()
     conf.override_all({})

def test_execute_failure_get_statement_output_emits_event():
    spark_events = MagicMock()
    kind = SESSION_KIND_SPARK
    http_client = MagicMock()
    conf.override_all({
        "status_sleep_seconds": 0.01,
        "statement_sleep_seconds": 0.01
    })
    session = _create_session(kind=kind, http_client=http_client)
    session.wait_for_idle = MagicMock()
    conf.override_all({})
    session.start(create_sql_context=False)
    session.wait_for_idle = MagicMock()
    command = Command("command", spark_events=spark_events)
    command._get_statement_output = MagicMock(side_effect=AttributeError('OHHHH'))

    try:
        result = command.execute(session)
        assert False
    except AttributeError as e:
        spark_events.emit_statement_execution_start_event._assert_called_once_with(session.guid, session.kind,
                                                                                   session.id, command.guid)
        spark_events.emit_statement_execution_start_event._assert_called_once_with(session.guid, session.kind,
                                                                                   session.id, command.guid,
                                                                                   -1, False, "AttributeError",
                                                                                   "OHHHH")
        assert_equals(e, command._get_statement_output.side_effect)

 def test_constructor_throws_invalid_session_sql_combo(self):
     conf.override_all({
         "status_sleep_seconds": 2,
         "statement_sleep_seconds": 2,
         "create_sql_context_timeout_seconds": 60
     })
     self._create_session(sql_created=True)
     conf.override_all({})

 def test_constructor_throws_status_sleep_seconds(self):
     conf.override_all({
         "status_sleep_seconds": 0,
         "statement_sleep_seconds": 2,
         "create_sql_context_timeout_seconds": 60
     })
     self._create_session()
     conf.override_all({})

def test_configuration_override_work_with_empty_password():
    kpc = { 'username': 'U', 'base64_password': '', 'password': '', 'url': '' }
    overrides = { conf.kernel_python_credentials.__name__: kpc }
    conf.override_all(overrides)
    conf.override(conf.status_sleep_seconds.__name__, 1)
    assert_equals(conf.d, { conf.kernel_python_credentials.__name__: kpc,
                                     conf.status_sleep_seconds.__name__: 1 })
    assert_equals(conf.status_sleep_seconds(), 1)
    assert_equals(conf.base64_kernel_python_credentials(),  { 'username': 'U', 'password': '', 'url': '' })

    def test_constructor_starts_with_no_session(self):
        conf.override_all({
            "status_sleep_seconds": 4,
            "statement_sleep_seconds": 2
        })
        session = self._create_session()
        conf.override_all({})

        assert session.id == -1

    def test_constructor_starts_with_no_session(self):
        conf.override_all({
            "status_sleep_seconds": 4,
            "statement_sleep_seconds": 2,
            "create_sql_context_timeout_seconds": 60
        })
        session = self._create_session()
        conf.override_all({})

        assert session.id == -1
        assert not session.created_sql_context

def _setup():
    global magic, spark_controller, shell, ipython_display, spark_events, conf

    conf.override_all({})
    spark_events = MagicMock()

    magic = TestKernelMagics(shell=None, spark_events=spark_events)
    magic.shell = shell = MagicMock()
    magic.ipython_display = ipython_display = MagicMock()
    magic.spark_controller = spark_controller = MagicMock()
    magic._generate_uuid = MagicMock(return_value='0000')

    def test_constructor_starts_with_existing_session(self):
        conf.override_all({
            "status_sleep_seconds": 4,
            "statement_sleep_seconds": 2
        })
        session_id = 1
        session = self._create_session(session_id=session_id, heartbeat_timeout=0)
        conf.override_all({})

        assert session.id == session_id
        assert session._heartbeat_thread is None
        assert constants.LIVY_HEARTBEAT_TIMEOUT_PARAM not in list(session.properties.keys())

    def test_delete_session_when_not_started(self):
        self.http_client.post_session.return_value = self.session_create_json
        conf.override_all({
            "status_sleep_seconds": 0.01,
            "statement_sleep_seconds": 0.01
        })
        session = self._create_session()
        conf.override_all({})

        session.delete()

        assert_equals(session.ipython_display.send_error.call_count, 1)

def test_sqlquery_loads_defaults():
    defaults = {
        conf.default_samplemethod.__name__: "sample",
        conf.default_maxrows.__name__: 419,
        conf.default_samplefraction.__name__: 0.99,
    }
    conf.override_all(defaults)
    query = "DROP TABLE USERS;"
    sqlquery = SQLQuery(query)
    assert_equals(sqlquery.query, query)
    assert_equals(sqlquery.samplemethod, defaults[conf.default_samplemethod.__name__])
    assert_equals(sqlquery.maxrows, defaults[conf.default_maxrows.__name__])
    assert_equals(sqlquery.samplefraction, defaults[conf.default_samplefraction.__name__])

    def test_is_final_status(self):
        conf.override_all({
            "status_sleep_seconds": 0.01,
            "statement_sleep_seconds": 0.01
        })
        session = self._create_session()
        conf.override_all({})

        assert not session.is_final_status("idle")
        assert not session.is_final_status("starting")
        assert not session.is_final_status("busy")

        assert session.is_final_status("dead")
        assert session.is_final_status("error")

    def test_constructor_starts_with_existing_session(self):
        conf.override_all({
            "status_sleep_seconds": 4,
            "statement_sleep_seconds": 2,
            "create_sql_context_timeout_seconds": 60
        })
        session_id = 1
        session = self._create_session(session_id=session_id, sql_created=True,
                                       should_heartbeat=False)
        conf.override_all({})

        assert session.id == session_id
        assert session.created_sql_context
        assert session._heartbeat_thread is None

 def test_create_sql_hive_context_unknown_throws(self):
     kind = "unknown"
     self.http_client.post_session.return_value = self.session_create_json
     self.http_client.post_statement.return_value = self.post_statement_json
     self.http_client.get_session.return_value = self.ready_sessions_json
     self.get_statement_responses = [self.running_statement_json, self.ready_statement_json]
     self.http_client.get_statement.side_effect = self._next_statement_response_get
     conf.override_all({
         "status_sleep_seconds": 0.01,
         "statement_sleep_seconds": 0.01
     })
     session = self._create_session(kind=kind)
     conf.override_all({})
     session.start()

    def test_delete_session_when_active(self):
        self.http_client.post_session.return_value = self.session_create_json
        self.http_client.get_session.return_value = self.ready_sessions_json
        conf.override_all({
            "status_sleep_seconds": 0.01,
            "statement_sleep_seconds": 0.01
        })
        session = self._create_session()
        conf.override_all({})
        session.start(create_sql_context=False)

        session.delete()

        assert_equals("dead", session.status)

 def test_constructor_starts_heartbeat_with_existing_session(self):
     conf.override_all({
         "status_sleep_seconds": 4,
         "statement_sleep_seconds": 2,
         "heartbeat_refresh_seconds": 0.1
     })
     session_id = 1
     session = self._create_session(session_id=session_id)
     conf.override_all({})
     
     assert session.id == session_id
     assert self.heartbeat_thread.daemon
     self.heartbeat_thread.start.assert_called_once_with()
     assert not session._heartbeat_thread is None
     assert session.properties[constants.LIVY_HEARTBEAT_TIMEOUT_PARAM ] > 0