C++ TestThriftServerFactory::create方法代码示例

TEST(ThriftServer, CompressionClientTest) {
  TestThriftServerFactory<TestInterface> factory;
  ScopedServerThread sst(factory.create());
  folly::EventBase base;
  std::shared_ptr<TAsyncSocket> socket(
      TAsyncSocket::newSocket(&base, *sst.getAddress()));

  TestServiceAsyncClient client(
      std::unique_ptr<HeaderClientChannel,
                      folly::DelayedDestruction::Destructor>(
          new HeaderClientChannel(socket)));

  auto channel =
      boost::polymorphic_downcast<HeaderClientChannel*>(client.getChannel());
  channel->setTransform(apache::thrift::transport::THeader::ZLIB_TRANSFORM);
  channel->setMinCompressBytes(1);

  std::string response;
  client.sync_sendResponse(response, 64);
  EXPECT_EQ(response, "test64");

  auto trans = channel->getWriteTransforms();
  EXPECT_EQ(trans.size(), 1);
  for (auto& tran : trans) {
    EXPECT_EQ(tran, apache::thrift::transport::THeader::ZLIB_TRANSFORM);
  }
}

TEST(ThriftServer, Thrift1OnewayRequestTest) {
  TestThriftServerFactory<TestInterface> factory;
  auto cpp2Server = factory.create();
  cpp2Server->setNWorkerThreads(1);
  cpp2Server->setIsOverloaded([](const THeader*) { return true; });
  apache::thrift::util::ScopedServerThread st(cpp2Server);

  std::shared_ptr<TestServiceClient> client =
      getThrift1Client(*st.getAddress());
  std::string response;
  // Send a oneway request. Server doesn't send error back
  client->noResponse(1);
  // Send a twoway request. Server sends overloading error back
  try {
    client->sendResponse(response, 0);
  } catch (apache::thrift::TApplicationException& ex) {
    EXPECT_STREQ(ex.what(), "loadshedding request");
  } catch (...) {
    ADD_FAILURE();
  }

  cpp2Server->setIsOverloaded([](const THeader*) { return false; });
  // Send another twoway request. Client should receive a response
  // with correct seqId
  client->sendResponse(response, 0);
}

TEST(ThriftServer, useExistingSocketAndConnectionIdleTimeout) {
  // This is ConnectionIdleTimeoutTest, but with an existing socket
  TestThriftServerFactory<TestInterface> factory;
  auto server = std::static_pointer_cast<ThriftServer>(factory.create());
  folly::AsyncServerSocket::UniquePtr serverSocket(
      new folly::AsyncServerSocket);
  serverSocket->bind(0);
  server->useExistingSocket(std::move(serverSocket));

  server->setIdleTimeout(std::chrono::milliseconds(20));
  apache::thrift::util::ScopedServerThread st(server);

  folly::EventBase base;
  std::shared_ptr<TAsyncSocket> socket(
      TAsyncSocket::newSocket(&base, *st.getAddress()));

  TestServiceAsyncClient client(
      std::unique_ptr<HeaderClientChannel,
                      folly::DelayedDestruction::Destructor>(
          new HeaderClientChannel(socket)));

  std::string response;
  client.sync_sendResponse(response, 200);
  EXPECT_EQ(response, "test200");
  base.loop();
}

TEST(ThriftServer, useExistingSocketAndExit) {
  TestThriftServerFactory<TestInterface> factory;
  auto server = std::static_pointer_cast<ThriftServer>(factory.create());
  folly::AsyncServerSocket::UniquePtr serverSocket(
      new folly::AsyncServerSocket);
  serverSocket->bind(0);
  server->useExistingSocket(std::move(serverSocket));
  // In the past, this would cause a SEGV
}

TEST(ThriftServer, IdleServerTimeout) {
  TestThriftServerFactory<TestInterface> factory;

  auto server = factory.create();
  auto thriftServer = dynamic_cast<ThriftServer *>(server.get());
  thriftServer->setIdleServerTimeout(std::chrono::milliseconds(50));

  ScopedServerThread scopedServer(server);
  scopedServer.join();
}

TEST(ThriftServer, setIOThreadPool) {
  auto exe = std::make_shared<wangle::IOThreadPoolExecutor>(1);
  TestThriftServerFactory<TestInterface> factory;
  factory.useSimpleThreadManager(false);
  auto server = std::static_pointer_cast<ThriftServer>(factory.create());

  // Set the exe, this used to trip various calls like
  // CHECK(ioThreadPool->numThreads() == 0).
  server->setIOThreadPool(exe);
  EXPECT_EQ(1, server->getNWorkerThreads());
}

TEST(ThriftServer, BadSendTest) {
  TestThriftServerFactory<TestInterface> factory;
  ScopedServerThread sst(factory.create());
  folly::EventBase base;
  std::shared_ptr<TAsyncSocket> socket(
      TAsyncSocket::newSocket(&base, *sst.getAddress()));

  TestServiceAsyncClient client(HeaderClientChannel::newChannel(socket));

  client.sendResponse(std::unique_ptr<RequestCallback>(new Callback), 64);

  socket->shutdownWriteNow();
  base.loop();

  std::string response;
  EXPECT_THROW(client.sync_sendResponse(response, 64), TTransportException);
}

TEST(ThriftServer, ConnectionIdleTimeoutTest) {
  TestThriftServerFactory<TestInterface> factory;
  auto server = factory.create();
  server->setIdleTimeout(std::chrono::milliseconds(20));
  apache::thrift::util::ScopedServerThread st(server);

  folly::EventBase base;
  std::shared_ptr<TAsyncSocket> socket(
      TAsyncSocket::newSocket(&base, *st.getAddress()));

  TestServiceAsyncClient client(HeaderClientChannel::newChannel(socket));

  std::string response;
  client.sync_sendResponse(response, 200);
  EXPECT_EQ(response, "test200");
  base.loop();
}

TEST(ThriftServer, FailureInjection) {
  enum ExpectedFailure { NONE = 0, ERROR, TIMEOUT, DISCONNECT, END };

  std::atomic<ExpectedFailure> expected(NONE);

  using apache::thrift::transport::TTransportException;

  class Callback : public RequestCallback {
   public:
    explicit Callback(const std::atomic<ExpectedFailure>* expected)
        : expected_(expected) {}

   private:
    void requestSent() override {}

    void replyReceived(ClientReceiveState&& state) override {
      std::string response;
      try {
        TestServiceAsyncClient::recv_sendResponse(response, state);
        EXPECT_EQ(NONE, *expected_);
      } catch (const apache::thrift::TApplicationException& ex) {
        const auto& headers = state.header()->getHeaders();
        EXPECT_TRUE(headers.find("ex") != headers.end() &&
                    headers.find("ex")->second == kInjectedFailureErrorCode);
        EXPECT_EQ(ERROR, *expected_);
      } catch (...) {
        ADD_FAILURE() << "Unexpected exception thrown";
      }

      // Now do it again with exception_wrappers.
      auto ew =
          TestServiceAsyncClient::recv_wrapped_sendResponse(response, state);
      if (ew) {
        EXPECT_TRUE(
            ew.is_compatible_with<apache::thrift::TApplicationException>());
        EXPECT_EQ(ERROR, *expected_);
      } else {
        EXPECT_EQ(NONE, *expected_);
      }
    }

    void requestError(ClientReceiveState&& state) override {
      try {
        std::rethrow_exception(state.exception());
      } catch (const TTransportException& ex) {
        if (ex.getType() == TTransportException::TIMED_OUT) {
          EXPECT_EQ(TIMEOUT, *expected_);
        } else {
          EXPECT_EQ(DISCONNECT, *expected_);
        }
      } catch (...) {
        ADD_FAILURE() << "Unexpected exception thrown";
      }
    }

    const std::atomic<ExpectedFailure>* expected_;
  };

  TestThriftServerFactory<TestInterface> factory;
  ScopedServerThread sst(factory.create());
  folly::EventBase base;
  std::shared_ptr<TAsyncSocket> socket(
      TAsyncSocket::newSocket(&base, *sst.getAddress()));

  TestServiceAsyncClient client(
      std::unique_ptr<HeaderClientChannel,
                      folly::DelayedDestruction::Destructor>(
          new HeaderClientChannel(socket)));

  auto server = std::dynamic_pointer_cast<ThriftServer>(sst.getServer().lock());
  CHECK(server);
  SCOPE_EXIT { server->setFailureInjection(ThriftServer::FailureInjection()); };

  RpcOptions rpcOptions;
  rpcOptions.setTimeout(std::chrono::milliseconds(100));
  for (int i = 0; i < END; ++i) {
    auto exp = static_cast<ExpectedFailure>(i);
    ThriftServer::FailureInjection fi;

    switch (exp) {
      case NONE:
        break;
      case ERROR:
        fi.errorFraction = 1;
        break;
      case TIMEOUT:
        fi.dropFraction = 1;
        break;
      case DISCONNECT:
        fi.disconnectFraction = 1;
        break;
      case END:
        LOG(FATAL) << "unreached";
        break;
    }

    server->setFailureInjection(std::move(fi));

    expected = exp;

//.........这里部分代码省略.........

TEST(ThriftServer, DefaultCompressionTest) {
  /* Tests the functionality of default transforms, ensuring the server properly
     applies them even if the client does not apply any transforms. */
  class Callback : public RequestCallback {
   public:
    explicit Callback(bool compressionExpected, uint16_t expectedTransform)
        : compressionExpected_(compressionExpected),
          expectedTransform_(expectedTransform) {}

   private:
    void requestSent() override {}

    void replyReceived(ClientReceiveState&& state) override {
      auto trans = state.header()->getTransforms();
      if (compressionExpected_) {
        EXPECT_EQ(trans.size(), 1);
        for (auto& tran : trans) {
          EXPECT_EQ(tran, expectedTransform_);
        }
      } else {
        EXPECT_EQ(trans.size(), 0);
      }
    }
    void requestError(ClientReceiveState&& state) override {
      std::rethrow_exception(state.exception());
    }
    bool compressionExpected_;
    uint16_t expectedTransform_;
  };

  TestThriftServerFactory<TestInterface> factory;
  factory.minCompressBytes(1);
  factory.defaultWriteTransform(
    apache::thrift::transport::THeader::ZLIB_TRANSFORM);
  auto server = std::static_pointer_cast<ThriftServer>(factory.create());
  ScopedServerThread sst(server);
  folly::EventBase base;

  // First, with minCompressBytes set low, ensure we compress even though the
  // client did not compress
  std::shared_ptr<TAsyncSocket> socket(
      TAsyncSocket::newSocket(&base, *sst.getAddress()));
  TestServiceAsyncClient client(HeaderClientChannel::newChannel(socket));
  client.sendResponse(
    folly::make_unique<Callback>(
      true, apache::thrift::transport::THeader::ZLIB_TRANSFORM
    ),
    64
  );
  base.loop();

  // Ensure that client transforms take precedence
  auto channel =
    boost::polymorphic_downcast<HeaderClientChannel*>(client.getChannel());
  channel->setTransform(apache::thrift::transport::THeader::SNAPPY_TRANSFORM);
  client.sendResponse(
    folly::make_unique<Callback>(
      true, apache::thrift::transport::THeader::SNAPPY_TRANSFORM
    ),
    64
  );
  base.loop();

  // Ensure that minCompressBytes still works with default transforms. We
  // Do not expect compression
  server->setMinCompressBytes(1000);
  std::shared_ptr<TAsyncSocket> socket2(
      TAsyncSocket::newSocket(&base, *sst.getAddress()));
  TestServiceAsyncClient client2(HeaderClientChannel::newChannel(socket2));
  client2.sendResponse(folly::make_unique<Callback>(false, 0), 64);
  base.loop();

}