PageRenderTime 20ms CodeModel.GetById 21ms RepoModel.GetById 0ms app.codeStats 0ms

/swift-service/src/test/java/com/facebook/swift/service/guice/TestThriftClientAndServerModules.java

https://gitlab.com/mayakarya/swift
Java | 312 lines | 243 code | 28 blank | 41 comment | 0 complexity | 2239da9adc1befe7d733c88b632a4e4d MD5 | raw file
  1. /*
  2. * Copyright (C) 2012 Facebook, Inc.
  3. *
  4. * Licensed under the Apache License, Version 2.0 (the "License"); you may
  5. * not use this file except in compliance with the License. You may obtain
  6. * a copy of the License at
  7. *
  8. * http://www.apache.org/licenses/LICENSE-2.0
  9. *
  10. * Unless required by applicable law or agreed to in writing, software
  11. * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  12. * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  13. * License for the specific language governing permissions and limitations
  14. * under the License.
  15. */
  16. package com.facebook.swift.service.guice;
  17. import com.facebook.nifty.client.FramedClientConnector;
  18. import com.facebook.nifty.client.NiftyClientChannel;
  19. import com.facebook.nifty.client.NiftyClientConnector;
  20. import com.facebook.nifty.core.RequestContext;
  21. import com.facebook.swift.codec.guice.ThriftCodecModule;
  22. import com.facebook.swift.service.LogEntry;
  23. import com.facebook.swift.service.ResultCode;
  24. import com.facebook.swift.service.Scribe;
  25. import com.facebook.swift.service.SwiftScribe;
  26. import com.facebook.swift.service.ThriftClient;
  27. import com.facebook.swift.service.ThriftClientConfig;
  28. import com.facebook.swift.service.ThriftEventHandler;
  29. import com.facebook.swift.service.ThriftServer;
  30. import com.facebook.swift.service.puma.TestPuma;
  31. import com.facebook.swift.service.puma.swift.PumaReadServer;
  32. import com.facebook.swift.service.puma.swift.PumaReadService;
  33. import com.facebook.swift.service.puma.swift.ReadResultQueryInfoTimeString;
  34. import com.google.common.collect.ImmutableList;
  35. import com.google.common.collect.ImmutableMap;
  36. import com.google.common.net.HostAndPort;
  37. import com.google.inject.Binder;
  38. import com.google.inject.BindingAnnotation;
  39. import com.google.inject.Guice;
  40. import com.google.inject.Injector;
  41. import com.google.inject.Key;
  42. import com.google.inject.Module;
  43. import com.google.inject.Scopes;
  44. import com.google.inject.Stage;
  45. import com.google.inject.TypeLiteral;
  46. import io.airlift.configuration.ConfigurationFactory;
  47. import io.airlift.configuration.ConfigurationModule;
  48. import io.airlift.jmx.testing.TestingJmxModule;
  49. import io.airlift.units.Duration;
  50. import org.testng.annotations.Test;
  51. import org.weakref.jmx.guice.MBeanModule;
  52. import java.lang.annotation.Retention;
  53. import java.lang.annotation.Target;
  54. import java.util.List;
  55. import static com.facebook.swift.service.guice.ThriftClientBinder.thriftClientBinder;
  56. import static com.facebook.swift.service.guice.ThriftServiceExporter.thriftServerBinder;
  57. import static com.facebook.swift.service.puma.TestPuma.verifyPumaResults;
  58. import static com.google.common.collect.Lists.newArrayList;
  59. import static java.lang.annotation.ElementType.CONSTRUCTOR;
  60. import static java.lang.annotation.ElementType.FIELD;
  61. import static java.lang.annotation.ElementType.METHOD;
  62. import static java.lang.annotation.ElementType.PARAMETER;
  63. import static java.lang.annotation.RetentionPolicy.RUNTIME;
  64. import static org.testng.Assert.assertEquals;
  65. public class TestThriftClientAndServerModules
  66. {
  67. private static final List<LogEntry> MESSAGES = ImmutableList.of(
  68. new LogEntry("hello", "world"),
  69. new LogEntry("bye", "world")
  70. );
  71. static private class TestThriftEventHandler extends ThriftEventHandler
  72. {
  73. public int count = 0;
  74. @Override
  75. public Object getContext(String methodName, RequestContext requestContext)
  76. {
  77. count++;
  78. return null;
  79. }
  80. }
  81. @Test
  82. public void testThriftClientAndServerModules()
  83. throws Exception
  84. {
  85. final List<Object> eventHandlerContexts = newArrayList();
  86. Injector injector = Guice.createInjector(Stage.PRODUCTION,
  87. new ConfigurationModule(new ConfigurationFactory(ImmutableMap.<String, String>of())),
  88. new ThriftCodecModule(),
  89. new ThriftClientModule(),
  90. new ThriftServerModule(),
  91. new TestingJmxModule(),
  92. new MBeanModule(),
  93. new Module()
  94. {
  95. @Override
  96. public void configure(Binder binder)
  97. {
  98. // bind scribe client
  99. thriftClientBinder(binder).bindThriftClient(Scribe.class);
  100. // bind scribe service implementation
  101. binder.bind(SwiftScribe.class).in(Scopes.SINGLETON);
  102. // export scribe service implementation
  103. thriftServerBinder(binder).exportThriftService(SwiftScribe.class);
  104. // bind puma client
  105. thriftClientBinder(binder).bindThriftClient(PumaReadService.class);
  106. // bind puma service implementation
  107. binder.bind(PumaReadServer.class).in(Scopes.SINGLETON);
  108. // export puma service implementation
  109. thriftServerBinder(binder).exportThriftService(PumaReadServer.class);
  110. // create an instance event handler
  111. thriftServerBinder(binder).addEventHandler(new ThriftEventHandler() {
  112. @Override
  113. public Object getContext(String methodName, RequestContext requestContext)
  114. {
  115. eventHandlerContexts.add(new Object());
  116. return null;
  117. }
  118. });
  119. // create a class event handler
  120. binder.bind(TestThriftEventHandler.class).in(Scopes.SINGLETON);
  121. thriftServerBinder(binder).addEventHandler(TestThriftEventHandler.class);
  122. }
  123. });
  124. try (ThriftServer server = injector.getInstance(ThriftServer.class).start()) {
  125. // test scribe
  126. ThriftClient<Scribe> scribeClient = injector.getInstance(Key.get(new TypeLiteral<ThriftClient<Scribe>>() {}));
  127. try (Scribe scribe = scribeClient.open(localFramedConnector(server.getPort())).get()) {
  128. assertEquals(scribe.log(MESSAGES), ResultCode.OK);
  129. assertEquals(injector.getInstance(SwiftScribe.class).getMessages(), newArrayList(MESSAGES));
  130. assertEquals(eventHandlerContexts.size(), 1);
  131. assertEquals(injector.getInstance(TestThriftEventHandler.class).count, 1);
  132. }
  133. // test puma
  134. ThriftClient<PumaReadService> pumaClient = injector.getInstance(Key.get(new TypeLiteral<ThriftClient<PumaReadService>>() {}));
  135. try (PumaReadService puma = pumaClient.open(localFramedConnector(server.getPort())).get()) {
  136. List<ReadResultQueryInfoTimeString> results = puma.getResultTimeString(TestPuma.PUMA_REQUEST);
  137. verifyPumaResults(results);
  138. assertEquals(eventHandlerContexts.size(), 2);
  139. assertEquals(injector.getInstance(TestThriftEventHandler.class).count, 2);
  140. }
  141. }
  142. }
  143. @Test
  144. public void testThriftWithAnnotationBinding()
  145. throws Exception
  146. {
  147. Injector injector = Guice.createInjector(Stage.PRODUCTION,
  148. new ConfigurationModule(new ConfigurationFactory(ImmutableMap.<String, String>of())),
  149. new ThriftCodecModule(),
  150. new ThriftClientModule(),
  151. new ThriftServerModule(),
  152. new Module()
  153. {
  154. @Override
  155. public void configure(Binder binder)
  156. {
  157. // bind scribe client
  158. thriftClientBinder(binder).bindThriftClient(Scribe.class, TestAnnotation.class);
  159. // bind scribe service implementation
  160. binder.bind(SwiftScribe.class).annotatedWith(TestAnnotation.class).to(SwiftScribe.class).in(Scopes.SINGLETON);
  161. // export scribe service implementation
  162. thriftServerBinder(binder).exportThriftService(SwiftScribe.class, TestAnnotation.class);
  163. // bind puma client
  164. thriftClientBinder(binder).bindThriftClient(PumaReadService.class, TestAnnotation.class);
  165. // bind puma service implementation
  166. binder.bind(PumaReadServer.class).annotatedWith(TestAnnotation.class).to(PumaReadServer.class).in(Scopes.SINGLETON);
  167. // export puma service implementation
  168. thriftServerBinder(binder).exportThriftService(PumaReadServer.class, TestAnnotation.class);
  169. }
  170. });
  171. try (ThriftServer server = injector.getInstance(ThriftServer.class).start()) {
  172. // test scribe
  173. ThriftClient<Scribe> scribeClient = injector.getInstance(Key.get(new TypeLiteral<ThriftClient<Scribe>>() {}, TestAnnotation.class));
  174. try (Scribe scribe = scribeClient.open(localFramedConnector(server.getPort())).get()) {
  175. assertEquals(scribe.log(MESSAGES), ResultCode.OK);
  176. assertEquals(injector.getInstance(Key.get(SwiftScribe.class, TestAnnotation.class)).getMessages(), newArrayList(MESSAGES));
  177. }
  178. // test puma
  179. ThriftClient<PumaReadService> pumaClient = injector.getInstance(Key.get(new TypeLiteral<ThriftClient<PumaReadService>>() {}, TestAnnotation.class));
  180. try (PumaReadService puma = pumaClient.open(localFramedConnector(server.getPort())).get()) {
  181. List<ReadResultQueryInfoTimeString> results = puma.getResultTimeString(TestPuma.PUMA_REQUEST);
  182. verifyPumaResults(results);
  183. }
  184. }
  185. }
  186. @Test
  187. public void testThriftWithKeyBinding()
  188. throws Exception
  189. {
  190. Injector injector = Guice.createInjector(Stage.PRODUCTION,
  191. new ConfigurationModule(new ConfigurationFactory(ImmutableMap.<String, String>of())),
  192. new ThriftCodecModule(),
  193. new ThriftClientModule(),
  194. new ThriftServerModule(),
  195. new Module()
  196. {
  197. @Override
  198. public void configure(Binder binder)
  199. {
  200. // bind scribe client
  201. thriftClientBinder(binder).bindThriftClient(Scribe.class);
  202. // bind scribe service implementation
  203. binder.bind(SwiftScribe.class).annotatedWith(TestAnnotation.class).to(SwiftScribe.class).in(Scopes.SINGLETON);
  204. // export scribe service implementation
  205. thriftServerBinder(binder).exportThriftService(Key.get(SwiftScribe.class, TestAnnotation.class));
  206. // bind puma client
  207. thriftClientBinder(binder).bindThriftClient(PumaReadService.class);
  208. // bind puma service implementation
  209. binder.bind(PumaReadServer.class).annotatedWith(TestAnnotation.class).to(PumaReadServer.class).in(Scopes.SINGLETON);
  210. // export puma service implementation
  211. thriftServerBinder(binder).exportThriftService(Key.get(PumaReadServer.class, TestAnnotation.class));
  212. }
  213. });
  214. try (ThriftServer server = injector.getInstance(ThriftServer.class).start()) {
  215. // test scribe
  216. ThriftClient<Scribe> scribeClient = injector.getInstance(Key.get(new TypeLiteral<ThriftClient<Scribe>>() {}));
  217. try (Scribe scribe = scribeClient.open(localFramedConnector(server.getPort())).get()) {
  218. assertEquals(scribe.log(MESSAGES), ResultCode.OK);
  219. assertEquals(injector.getInstance(Key.get(SwiftScribe.class, TestAnnotation.class)).getMessages(), newArrayList(MESSAGES));
  220. }
  221. // test puma
  222. ThriftClient<PumaReadService> pumaClient = injector.getInstance(Key.get(new TypeLiteral<ThriftClient<PumaReadService>>() {}));
  223. try (PumaReadService puma = pumaClient.open(localFramedConnector(server.getPort())).get()) {
  224. List<ReadResultQueryInfoTimeString> results = puma.getResultTimeString(TestPuma.PUMA_REQUEST);
  225. verifyPumaResults(results);
  226. }
  227. }
  228. }
  229. @Test
  230. public void testThriftClientWithConfiguration()
  231. throws Exception
  232. {
  233. final String SOCKS_PROXY_HOSTNAME = "proxyserver";
  234. final int SOCKS_PROXY_PORT = 1080;
  235. final int SIXTEEN_MB_IN_BYTES = 16777216;
  236. final String TEST_MEDIUM_TIMEOUT = "1s";
  237. final String TEST_SHORT_TIMEOUT = "150ms";
  238. final String TEST_LONG_TIMEOUT = "1m";
  239. HostAndPort proxy = HostAndPort.fromParts(SOCKS_PROXY_HOSTNAME, SOCKS_PROXY_PORT);
  240. ImmutableMap<String, String> configMap = new ImmutableMap.Builder<String, String>()
  241. .put("scribe.thrift.client.connect-timeout", TEST_MEDIUM_TIMEOUT)
  242. .put("scribe.thrift.client.receive-timeout", TEST_MEDIUM_TIMEOUT)
  243. .put("scribe.thrift.client.read-timeout", TEST_SHORT_TIMEOUT)
  244. .put("scribe.thrift.client.max-frame-size", Integer.toString(SIXTEEN_MB_IN_BYTES))
  245. .put("PumaReadService.thrift.client.write-timeout", TEST_LONG_TIMEOUT)
  246. .put("PumaReadService.thrift.client.socks-proxy", proxy.toString())
  247. .build();
  248. Injector injector = Guice.createInjector(
  249. Stage.PRODUCTION,
  250. new ConfigurationModule(new ConfigurationFactory(configMap)),
  251. new ThriftCodecModule(),
  252. new ThriftClientModule(),
  253. new Module() {
  254. @Override
  255. public void configure(Binder binder)
  256. {
  257. thriftClientBinder(binder).bindThriftClient(Scribe.class);
  258. thriftClientBinder(binder).bindThriftClient(PumaReadService.class);
  259. }
  260. });
  261. ThriftClient<Scribe> scribeClient = injector.getInstance(Key.get(new TypeLiteral<ThriftClient<Scribe>>() {}));
  262. assertEquals(Duration.valueOf(scribeClient.getConnectTimeout()), Duration.valueOf(TEST_MEDIUM_TIMEOUT));
  263. assertEquals(Duration.valueOf(scribeClient.getReceiveTimeout()), Duration.valueOf(TEST_MEDIUM_TIMEOUT));
  264. assertEquals(Duration.valueOf(scribeClient.getReadTimeout()), Duration.valueOf(TEST_SHORT_TIMEOUT));
  265. assertEquals(Duration.valueOf(scribeClient.getWriteTimeout()), ThriftClientConfig.DEFAULT_WRITE_TIMEOUT);
  266. assertEquals(scribeClient.getMaxFrameSize(), SIXTEEN_MB_IN_BYTES);
  267. assertEquals(scribeClient.getSocksProxy(), null);
  268. ThriftClient<PumaReadService> pumaClient = injector.getInstance(Key.get(new TypeLiteral<ThriftClient<PumaReadService>>() {}));
  269. assertEquals(Duration.valueOf(pumaClient.getConnectTimeout()), ThriftClientConfig.DEFAULT_CONNECT_TIMEOUT);
  270. assertEquals(Duration.valueOf(pumaClient.getReadTimeout()), ThriftClientConfig.DEFAULT_READ_TIMEOUT);
  271. assertEquals(Duration.valueOf(pumaClient.getReceiveTimeout()), ThriftClientConfig.DEFAULT_RECEIVE_TIMEOUT);
  272. assertEquals(Duration.valueOf(pumaClient.getWriteTimeout()), Duration.valueOf(TEST_LONG_TIMEOUT));
  273. assertEquals(pumaClient.getMaxFrameSize(), ThriftClientConfig.DEFAULT_MAX_FRAME_SIZE);
  274. assertEquals(HostAndPort.fromString(pumaClient.getSocksProxy()), proxy);
  275. }
  276. private NiftyClientConnector<? extends NiftyClientChannel> localFramedConnector(int port) {
  277. return new FramedClientConnector(HostAndPort.fromParts("localhost", port));
  278. }
  279. @Target({METHOD, CONSTRUCTOR, FIELD, PARAMETER})
  280. @Retention(RUNTIME)
  281. @BindingAnnotation
  282. public @interface TestAnnotation {
  283. }
  284. }