STRLCPY/tsunami-security-scanner

Add HealthCheck gRPC client implementation.

PiperOrigin-RevId: 460548451
Change-Id: I5dc93fbe1052521bebbe23d6c3e6430c060d5616

John Y. Kim committed with Copybara-Service 2 years ago

524cba70

1 parent 59fcb57b

Revision indexing in progress... (symbol navigation in revisions will be accurate after indexed)

Total 4 files Show one by one

■ ■ ■ ■ ■ ■

build.gradle

		skipped 64 lines
65	65		grpc_context: "io.grpc:grpc-context:${grpcVersion}",
66	66		grpc_stub: "io.grpc:grpc-stub:${grpcVersion}",
67	67		grpc_core: "io.grpc:grpc-core:${grpcVersion}",
	68	+	grpc_services: "io.grpc:grpc-services:${grpcVersion}",
68	69		grpc_testing: "io.grpc:grpc-testing:${grpcVersion}",
69	70		javax_annotations: "javax.annotation:javax.annotation-api:${javaxAnnotationVersion}",
70	71		gson: "com.google.code.gson:gson:${gsonVersion}",
		skipped 155 lines

■ ■ ■ ■ ■ ■

plugin/build.gradle

		skipped 16 lines
17	17		compile deps.grpc_core
18	18		compile deps.grpc_testing
19	19		compile deps.grpc_context
	20	+	compile deps.grpc_services
20	21		compile deps.tcs_common, deps.tcs_proto
21	22		annotationProcessor deps.autovalue_annotation_processor
22	23
		skipped 6 lines

■ ■ ■ ■ ■ ■

plugin/src/main/java/com/google/tsunami/plugin/PluginServiceClient.java

		skipped 25 lines
26	26		import com.google.tsunami.proto.RunResponse;
27	27		import io.grpc.Channel;
28	28		import io.grpc.Deadline;
	29	+	import io.grpc.health.v1.HealthCheckRequest;
	30	+	import io.grpc.health.v1.HealthCheckResponse;
	31	+	import io.grpc.health.v1.HealthGrpc;
	32	+	import io.grpc.health.v1.HealthGrpc.HealthFutureStub;
29	33
30	34		/**
31	35		* Client side gRPC handler for the PluginService RPC protocol. Main handler for all gRPC calls to
		skipped 1 lines
33	37		*/
34	38		public final class PluginServiceClient {
35	39
	40	+	private final HealthFutureStub healthService;
36	41		private final PluginServiceFutureStub pluginService;
37	42
38	43		PluginServiceClient(Channel channel) {
	44	+	this.healthService = HealthGrpc.newFutureStub(checkNotNull(channel));
39	45		this.pluginService = PluginServiceGrpc.newFutureStub(checkNotNull(channel));
40	46		}
41	47
		skipped 18 lines
60	66		public ListenableFuture<ListPluginsResponse> listPluginsWithDeadline(
61	67		ListPluginsRequest request, Deadline deadline) {
62	68		return pluginService.withDeadline(deadline).listPlugins(request);
	69	+	}
	70	+
	71	+	/**
	72	+	* Sends a health check request to retrieve the status of the language server.
	73	+	*
	74	+	* @param request The health check request to send to the language server.
	75	+	* @param deadline The maximum time to keep this call alive.
	76	+	* @return The language server's status via {@link HealthCheckResponse}.
	77	+	*/
	78	+	public ListenableFuture<HealthCheckResponse> checkHealthWithDeadline(
	79	+	HealthCheckRequest request, Deadline deadline) {
	80	+	return healthService.withDeadline(deadline).check(request);
63	81		}
64	82		}
65	83

■ ■ ■ ■ ■ ■

plugin/src/test/java/com/google/tsunami/plugin/PluginServiceClientTest.java

		skipped 35 lines
36	36		import com.google.tsunami.proto.TargetInfo;
37	37		import com.google.tsunami.proto.TransportProtocol;
38	38		import io.grpc.Deadline;
	39	+	import io.grpc.health.v1.HealthCheckRequest;
	40	+	import io.grpc.health.v1.HealthCheckResponse;
	41	+	import io.grpc.health.v1.HealthCheckResponse.ServingStatus;
	42	+	import io.grpc.health.v1.HealthGrpc.HealthImplBase;
39	43		import io.grpc.inprocess.InProcessChannelBuilder;
40	44		import io.grpc.inprocess.InProcessServerBuilder;
41	45		import io.grpc.stub.StreamObserver;
		skipped 179 lines
221	225
222	226		assertThat(listPlugins.isDone()).isTrue();
223	227		assertThat(listPlugins.get().getPluginsList()).containsExactlyElementsIn(plugins);
	228	+	}
	229	+
	230	+	@Test
	231	+	public void checkHealth_returnServingHealthResponse() throws Exception {
	232	+	HealthCheckRequest request = HealthCheckRequest.getDefaultInstance();
	233	+
	234	+	HealthImplBase healthImpl =
	235	+	new HealthImplBase() {
	236	+	@Override
	237	+	public void check(
	238	+	HealthCheckRequest request, StreamObserver<HealthCheckResponse> responseObserver) {
	239	+	responseObserver.onNext(
	240	+	HealthCheckResponse.newBuilder().setStatus(ServingStatus.SERVING).build());
	241	+	responseObserver.onCompleted();
	242	+	}
	243	+	};
	244	+	serviceRegistry.addService(healthImpl);
	245	+
	246	+	ListenableFuture<HealthCheckResponse> health =
	247	+	pluginService.checkHealthWithDeadline(request, DEADLINE_DEFAULT);
	248	+
	249	+	assertThat(health.isDone()).isTrue();
	250	+	assertThat(health.get().getStatus()).isEqualTo(ServingStatus.SERVING);
	251	+	}
	252	+
	253	+	@Test
	254	+	public void checkHealth_returnNotServingHealthResponse() throws Exception {
	255	+	HealthCheckRequest request = HealthCheckRequest.getDefaultInstance();
	256	+
	257	+	HealthImplBase healthImpl =
	258	+	new HealthImplBase() {
	259	+	@Override
	260	+	public void check(
	261	+	HealthCheckRequest request, StreamObserver<HealthCheckResponse> responseObserver) {
	262	+	responseObserver.onNext(
	263	+	HealthCheckResponse.newBuilder().setStatus(ServingStatus.NOT_SERVING).build());
	264	+	responseObserver.onCompleted();
	265	+	}
	266	+	};
	267	+	serviceRegistry.addService(healthImpl);
	268	+
	269	+	ListenableFuture<HealthCheckResponse> health =
	270	+	pluginService.checkHealthWithDeadline(request, DEADLINE_DEFAULT);
	271	+
	272	+	assertThat(health.isDone()).isTrue();
	273	+	assertThat(health.get().getStatus()).isEqualTo(ServingStatus.NOT_SERVING);
224	274		}
225	275
226	276		private void assertRunResponseContainsAllRunRequestParameters(
		skipped 41 lines