Examples with CopyDoneMessage com.google.cloud.spanner.pgadapter.wireprotocol.CopyDoneMessage used on opensource projects

Example 1 with CopyDoneMessage

use of com.google.cloud.spanner.pgadapter.wireprotocol.CopyDoneMessage in project pgadapter by GoogleCloudPlatform.

the class ProtocolTest method testCopyDoneMessage.

@Test
public void testCopyDoneMessage() throws Exception {
    byte[] messageMetadata = { 'c' };
    byte[] length = intToBytes(4);
    byte[] value = Bytes.concat(messageMetadata, length);
    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
    ByteArrayOutputStream result = new ByteArrayOutputStream();
    DataOutputStream outputStream = new DataOutputStream(result);
    CopyStatement copyStatement = mock(CopyStatement.class);
    when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
    MutationWriter mb = mock(MutationWriter.class);
    when(copyStatement.getMutationWriter()).thenReturn(mb);
    WireMessage message = ControlMessage.create(connectionHandler);
    assertEquals(CopyDoneMessage.class, message.getClass());
    CopyDoneMessage messageSpy = (CopyDoneMessage) spy(message);
    doReturn(false).when(messageSpy).sendSpannerResult(anyInt(), any(IntermediateStatement.class), any(QueryMode.class), anyLong());
    messageSpy.send();
    verify(messageSpy).sendSpannerResult(0, copyStatement, QueryMode.SIMPLE, 0L);
}