### Question:
MessageInputStream extends DataInputStream { public Message readMessage() throws IOException { short id = readShort(); MessageType messageType = messageRegistry.getMessageType(id); if (messageType == null) { throw new MessageTypeException(id); } LinkedList<Class> classes = new LinkedList<>(); for (ArgumentType argumentType : messageType.getArguments()) { Class typeForName = Message.getTypeForName(argumentType.getType()); classes.add(typeForName); } Class<?>[] classArray = new Class[classes.size()]; classes.toArray(classArray); Object[] args = readMessageArgs(classes); return new Message(messageType, args); } MessageInputStream(InputStream inputStream); Message readMessage(); void setPrivateKey(Key privateKey);  }### Answer:
@Test public void testReadValidMessage() throws IOException { DataInputStreamMock dataInputStreamMock = new DataInputStreamMock(); dataInputStreamMock.mockReadShort((short) 0); dataInputStreamMock.mockReadBoolean(false); dataInputStreamMock.mockReadUtf("FooBar"); dataInputStreamMock.mockReadShort((short) 0); InputStream inputStream = dataInputStreamMock.getIn(); MessageInputStream messageInputStream = new MessageInputStream(inputStream); Message message = messageInputStream.readMessage(); messageInputStream.close(); Assert.assertEquals((String) message.getArgument("messageName"), "FooBar", "Argument names do not match."); }
@Test(expectedExceptions = MessageTypeException.class) public void testReadInvalidMessage() throws IOException { DataInputStreamMock mock = new DataInputStreamMock(); mock.mockReadShort((short) -256); mock.mockReadBoolean(false); mock.mockReadUtf("FooBar"); mock.mockReadShort((short) -1); InputStream is = mock.getIn(); MessageInputStream mis = new MessageInputStream(is); mis.readMessage(); }
@Test public void testEncryptedMessageNoKey() throws IOException { DataInputStreamMock mock = new DataInputStreamMock(); mock.mockReadShort((short) -3); mock.mockReadBoolean(true); mock.mockReadUtf("FooBar"); InputStream is = mock.getIn(); MessageInputStream mis = new MessageInputStream(is); Message msg = mis.readMessage(); mis.close(); Assert.assertEquals(msg.getArgument("encryptedString"), "FooBar", "Arguments do not match."); }
@Test public void testAllTypesMessage() throws Exception { DataInputStreamMock mock = new DataInputStreamMock(); mock.mockReadShort((short) -4); mock.mockReadBoolean(false); mock.mockReadUtf("FooBar"); mock.mockReadInt(123); mock.mockReadLong(456); mock.mockReadDouble(Math.random()); mock.mockReadFloat((float) Math.random()); mock.mockReadShort((short) 246); mock.mockReadBoolean(true); mock.mockReadByte((byte) 867); mock.mockReadInt(3); mock.mockReadByte((byte) 8); mock.mockReadByte((byte) 16); mock.mockReadByte((byte) 32); mock.mockReadInt(3); mock.mockReadBoolean(false); mock.mockReadUtf("I'm"); mock.mockReadBoolean(false); mock.mockReadUtf("A"); mock.mockReadBoolean(false); mock.mockReadUtf("Lumberjack"); InputStream in = mock.getIn(); MessageInputStream mis = new MessageInputStream(in); Message msg = mis.readMessage(); mis.close(); Assert.assertEquals(((String[]) msg.getArgument("testStringArray"))[1], "A", "String array contents incorrect"); }