1 /* 2 * Copyright 2018 The gRPC Authors 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain 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, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package io.grpc.alts.internal; 18 19 import static com.google.common.truth.Truth.assertThat; 20 import static org.junit.Assert.assertEquals; 21 import static org.junit.Assert.fail; 22 23 import com.google.protobuf.ByteString; 24 import io.grpc.alts.internal.Handshaker.HandshakerReq; 25 import io.grpc.alts.internal.Handshaker.HandshakerResp; 26 import io.grpc.alts.internal.Handshaker.NextHandshakeMessageReq; 27 import io.grpc.stub.StreamObserver; 28 import java.io.IOException; 29 import org.junit.Before; 30 import org.junit.Test; 31 import org.junit.runner.RunWith; 32 import org.junit.runners.JUnit4; 33 34 /** Unit tests for {@link AltsHandshakerStub}. */ 35 @RunWith(JUnit4.class) 36 public class AltsHandshakerStubTest { 37 /** Mock status of handshaker service. */ 38 private static enum Status { 39 OK, 40 ERROR, 41 COMPLETE 42 } 43 44 private AltsHandshakerStub stub; 45 private MockWriter writer; 46 47 @Before setUp()48 public void setUp() { 49 writer = new MockWriter(); 50 stub = new AltsHandshakerStub(writer); 51 writer.setReader(stub.getReaderForTest()); 52 } 53 54 /** Send a message as in_bytes and expect same message as out_frames echo back. */ sendSuccessfulMessage()55 private void sendSuccessfulMessage() throws Exception { 56 String message = "hello world"; 57 HandshakerReq.Builder req = 58 HandshakerReq.newBuilder() 59 .setNext( 60 NextHandshakeMessageReq.newBuilder() 61 .setInBytes(ByteString.copyFromUtf8(message)) 62 .build()); 63 HandshakerResp resp = stub.send(req.build()); 64 assertEquals(resp.getOutFrames().toStringUtf8(), message); 65 } 66 67 /** Send a message and expect an IOException on error. */ sendAndExpectError()68 private void sendAndExpectError() throws InterruptedException { 69 try { 70 stub.send(HandshakerReq.newBuilder().build()); 71 fail("Exception expected"); 72 } catch (IOException ex) { 73 assertThat(ex).hasMessageThat().contains("Received a terminating error"); 74 } 75 } 76 77 /** Send a message and expect an IOException on closing. */ sendAndExpectComplete()78 private void sendAndExpectComplete() throws InterruptedException { 79 try { 80 stub.send(HandshakerReq.newBuilder().build()); 81 fail("Exception expected"); 82 } catch (IOException ex) { 83 assertThat(ex).hasMessageThat().contains("Response stream closed"); 84 } 85 } 86 87 /** Send a message and expect an IOException on unexpected message. */ sendAndExpectUnexpectedMessage()88 private void sendAndExpectUnexpectedMessage() throws InterruptedException { 89 try { 90 stub.send(HandshakerReq.newBuilder().build()); 91 fail("Exception expected"); 92 } catch (IOException ex) { 93 assertThat(ex).hasMessageThat().contains("Received an unexpected response"); 94 } 95 } 96 97 @Test sendSuccessfulMessageTest()98 public void sendSuccessfulMessageTest() throws Exception { 99 writer.setServiceStatus(Status.OK); 100 sendSuccessfulMessage(); 101 stub.close(); 102 } 103 104 @Test getServiceErrorTest()105 public void getServiceErrorTest() throws InterruptedException { 106 writer.setServiceStatus(Status.ERROR); 107 sendAndExpectError(); 108 stub.close(); 109 } 110 111 @Test getServiceCompleteTest()112 public void getServiceCompleteTest() throws Exception { 113 writer.setServiceStatus(Status.COMPLETE); 114 sendAndExpectComplete(); 115 stub.close(); 116 } 117 118 @Test getUnexpectedMessageTest()119 public void getUnexpectedMessageTest() throws Exception { 120 writer.setServiceStatus(Status.OK); 121 writer.sendUnexpectedResponse(); 122 sendAndExpectUnexpectedMessage(); 123 stub.close(); 124 } 125 126 @Test closeEarlyTest()127 public void closeEarlyTest() throws InterruptedException { 128 stub.close(); 129 sendAndExpectComplete(); 130 } 131 132 private static class MockWriter implements StreamObserver<HandshakerReq> { 133 private StreamObserver<HandshakerResp> reader; 134 private Status status = Status.OK; 135 setReader(StreamObserver<HandshakerResp> reader)136 private void setReader(StreamObserver<HandshakerResp> reader) { 137 this.reader = reader; 138 } 139 setServiceStatus(Status status)140 private void setServiceStatus(Status status) { 141 this.status = status; 142 } 143 144 /** Send a handshaker response to reader. */ sendUnexpectedResponse()145 private void sendUnexpectedResponse() { 146 reader.onNext(HandshakerResp.newBuilder().build()); 147 } 148 149 /** Mock writer onNext. Will respond based on the server status. */ 150 @Override onNext(final HandshakerReq req)151 public void onNext(final HandshakerReq req) { 152 switch (status) { 153 case OK: 154 HandshakerResp.Builder resp = HandshakerResp.newBuilder(); 155 reader.onNext(resp.setOutFrames(req.getNext().getInBytes()).build()); 156 break; 157 case ERROR: 158 reader.onError(new RuntimeException()); 159 break; 160 case COMPLETE: 161 reader.onCompleted(); 162 break; 163 default: 164 return; 165 } 166 } 167 168 @Override onError(Throwable t)169 public void onError(Throwable t) {} 170 171 /** Mock writer onComplete. */ 172 @Override onCompleted()173 public void onCompleted() { 174 reader.onCompleted(); 175 } 176 } 177 } 178