diff --git a/src/facade/dragonfly_connection.cc b/src/facade/dragonfly_connection.cc index 367423c00..21849e476 100644 --- a/src/facade/dragonfly_connection.cc +++ b/src/facade/dragonfly_connection.cc @@ -698,6 +698,10 @@ auto Connection::FromArgs(RespVec args, mi_heap_t* heap) -> Request* { return req; } +void Connection::ShutdownSelf() { + util::Connection::Shutdown(); +} + void RespToArgList(const RespVec& src, CmdArgVec* dest) { dest->resize(src.size()); for (size_t i = 0; i < src.size(); ++i) { diff --git a/src/facade/dragonfly_connection.h b/src/facade/dragonfly_connection.h index 5578d5ed8..f7258af59 100644 --- a/src/facade/dragonfly_connection.h +++ b/src/facade/dragonfly_connection.h @@ -72,6 +72,8 @@ class Connection : public util::Connection { std::string GetClientInfo() const; uint32 GetClientId() const; + void ShutdownSelf(); + protected: void OnShutdown() override; void OnPreMigrateThread() override; diff --git a/src/server/main_service.cc b/src/server/main_service.cc index 49f1ad7cd..a8efec664 100644 --- a/src/server/main_service.cc +++ b/src/server/main_service.cc @@ -721,6 +721,7 @@ void Service::Quit(CmdArgList args, ConnectionContext* cntx) { SinkReplyBuilder* builder = cntx->reply_builder(); builder->CloseConnection(); + cntx->owner()->ShutdownSelf(); } void Service::Multi(CmdArgList args, ConnectionContext* cntx) { diff --git a/tests/dragonfly/conftest.py b/tests/dragonfly/conftest.py index 46473b1b1..08fca369b 100644 --- a/tests/dragonfly/conftest.py +++ b/tests/dragonfly/conftest.py @@ -48,7 +48,7 @@ def df_server(request, tmp_dir: Path, test_env): arguments = [arg.format(**test_env) for arg in request.param] dfly_proc = subprocess.Popen([DRAGONFLY_PATH, *arguments], env=test_env, cwd=str(tmp_dir)) - time.sleep(0.1) + time.sleep(0.3) return_code = dfly_proc.poll() if return_code is not None: dfly_proc.terminate() @@ -66,17 +66,19 @@ def df_server(request, tmp_dir: Path, test_env): print(outs) print(errs) +@pytest.fixture(scope="function") +def connection(df_server): + return redis.Connection() @pytest.fixture(scope="class") -def connection(df_server): +def raw_client(df_server): """ Creates the Redis client to interact with the Dragonfly instance """ pool = redis.ConnectionPool(decode_responses=True) client = redis.Redis(connection_pool=pool) return client - @pytest.fixture -def client(connection): +def client(raw_client): """ Flushes all the records, runs before each test. """ - connection.flushall() - return connection + raw_client.flushall() + return raw_client diff --git a/tests/dragonfly/server_family_test.py b/tests/dragonfly/server_family_test.py new file mode 100644 index 000000000..1d72d2a66 --- /dev/null +++ b/tests/dragonfly/server_family_test.py @@ -0,0 +1,20 @@ +import pytest +import redis + +def test_quit(connection): + connection.send_command("QUIT") + assert connection.read_response() == b'OK' + + with pytest.raises(redis.exceptions.ConnectionError) as e: + connection.read_response() + +def test_quit_after_sub(connection): + connection = redis.Connection() + connection.send_command("SUBSCRIBE", "foo") + connection.read_response() + + connection.send_command("QUIT") + assert connection.read_response() == b'OK' + + with pytest.raises(redis.exceptions.ConnectionError) as e: + connection.read_response()