Java tutorial
/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.flink.streaming.api.functions.source; import org.apache.commons.io.IOUtils; import org.apache.flink.streaming.api.watermark.Watermark; import org.junit.Test; import java.io.EOFException; import java.io.OutputStreamWriter; import java.net.ServerSocket; import java.net.Socket; import static org.junit.Assert.*; /** * Tests for the {@link org.apache.flink.streaming.api.functions.source.SocketTextStreamFunction}. */ public class SocketTextStreamFunctionTest { private static final String LOCALHOST = "127.0.0.1"; @Test public void testSocketSourceSimpleOutput() throws Exception { ServerSocket server = new ServerSocket(0); Socket channel = null; try { SocketTextStreamFunction source = new SocketTextStreamFunction(LOCALHOST, server.getLocalPort(), "\n", 0); SocketSourceThread runner = new SocketSourceThread(source, "test1", "check"); runner.start(); channel = server.accept(); OutputStreamWriter writer = new OutputStreamWriter(channel.getOutputStream()); writer.write("test1\n"); writer.write("check\n"); writer.flush(); runner.waitForNumElements(2); runner.cancel(); runner.interrupt(); runner.waitUntilDone(); channel.close(); } finally { if (channel != null) { IOUtils.closeQuietly(channel); } IOUtils.closeQuietly(server); } } @Test public void testExitNoRetries() throws Exception { ServerSocket server = new ServerSocket(0); Socket channel = null; try { SocketTextStreamFunction source = new SocketTextStreamFunction(LOCALHOST, server.getLocalPort(), "\n", 0); SocketSourceThread runner = new SocketSourceThread(source); runner.start(); channel = server.accept(); channel.close(); try { runner.waitUntilDone(); } catch (Exception e) { assertTrue(e.getCause() instanceof EOFException); } } finally { if (channel != null) { IOUtils.closeQuietly(channel); } IOUtils.closeQuietly(server); } } @Test public void testSocketSourceOutputWithRetries() throws Exception { ServerSocket server = new ServerSocket(0); Socket channel = null; try { SocketTextStreamFunction source = new SocketTextStreamFunction(LOCALHOST, server.getLocalPort(), "\n", 10, 100); SocketSourceThread runner = new SocketSourceThread(source, "test1", "check"); runner.start(); // first connection: nothing channel = server.accept(); channel.close(); // second connection: first string channel = server.accept(); OutputStreamWriter writer = new OutputStreamWriter(channel.getOutputStream()); writer.write("test1\n"); writer.close(); channel.close(); // third connection: nothing channel = server.accept(); channel.close(); // forth connection: second string channel = server.accept(); writer = new OutputStreamWriter(channel.getOutputStream()); writer.write("check\n"); writer.flush(); runner.waitForNumElements(2); runner.cancel(); runner.waitUntilDone(); } finally { if (channel != null) { IOUtils.closeQuietly(channel); } IOUtils.closeQuietly(server); } } @Test public void testSocketSourceOutputInfiniteRetries() throws Exception { ServerSocket server = new ServerSocket(0); Socket channel = null; try { SocketTextStreamFunction source = new SocketTextStreamFunction(LOCALHOST, server.getLocalPort(), "\n", -1, 100); SocketSourceThread runner = new SocketSourceThread(source, "test1", "check"); runner.start(); // first connection: nothing channel = server.accept(); channel.close(); // second connection: first string channel = server.accept(); OutputStreamWriter writer = new OutputStreamWriter(channel.getOutputStream()); writer.write("test1\n"); writer.close(); channel.close(); // third connection: nothing channel = server.accept(); channel.close(); // forth connection: second string channel = server.accept(); writer = new OutputStreamWriter(channel.getOutputStream()); writer.write("check\n"); writer.flush(); runner.waitForNumElements(2); runner.cancel(); runner.waitUntilDone(); } finally { if (channel != null) { IOUtils.closeQuietly(channel); } IOUtils.closeQuietly(server); } } @Test public void testSocketSourceOutputAcrossRetries() throws Exception { ServerSocket server = new ServerSocket(0); Socket channel = null; try { SocketTextStreamFunction source = new SocketTextStreamFunction(LOCALHOST, server.getLocalPort(), "\n", 10, 100); SocketSourceThread runner = new SocketSourceThread(source, "test1", "check1", "check2"); runner.start(); // first connection: nothing channel = server.accept(); channel.close(); // second connection: first string channel = server.accept(); OutputStreamWriter writer = new OutputStreamWriter(channel.getOutputStream()); writer.write("te"); writer.close(); channel.close(); // third connection: nothing channel = server.accept(); channel.close(); // forth connection: second string channel = server.accept(); writer = new OutputStreamWriter(channel.getOutputStream()); writer.write("st1\n"); writer.write("check1\n"); writer.write("check2\n"); writer.flush(); runner.waitForNumElements(2); runner.cancel(); runner.waitUntilDone(); } finally { if (channel != null) { IOUtils.closeQuietly(channel); } IOUtils.closeQuietly(server); } } // ------------------------------------------------------------------------ private static class SocketSourceThread extends Thread { private final Object sync = new Object(); private final SocketTextStreamFunction socketSource; private final String[] expectedData; private volatile Throwable error; private volatile int numElementsReceived; private volatile boolean canceled; private volatile boolean done; public SocketSourceThread(SocketTextStreamFunction socketSource, String... expectedData) { this.socketSource = socketSource; this.expectedData = expectedData; } public void run() { try { SourceFunction.SourceContext<String> ctx = new SourceFunction.SourceContext<String>() { private final Object lock = new Object(); @Override public void collect(String element) { int pos = numElementsReceived; // make sure waiter know of us synchronized (sync) { numElementsReceived++; sync.notifyAll(); } if (expectedData != null && expectedData.length > pos) { assertEquals(expectedData[pos], element); } } @Override public void collectWithTimestamp(String element, long timestamp) { collect(element); } @Override public void emitWatermark(Watermark mark) { throw new UnsupportedOperationException(); } @Override public void markAsTemporarilyIdle() { throw new UnsupportedOperationException(); } @Override public Object getCheckpointLock() { return lock; } @Override public void close() { } }; socketSource.run(ctx); } catch (Throwable t) { synchronized (sync) { if (!canceled) { error = t; } sync.notifyAll(); } } finally { synchronized (sync) { done = true; sync.notifyAll(); } } } public void cancel() { synchronized (sync) { canceled = true; socketSource.cancel(); interrupt(); } } public void waitForNumElements(int numElements) throws InterruptedException { synchronized (sync) { while (error == null && !canceled && !done && numElementsReceived < numElements) { sync.wait(); } if (error != null) { throw new RuntimeException("Error in source thread", error); } if (canceled) { throw new RuntimeException("canceled"); } if (done) { throw new RuntimeException("Exited cleanly before expected number of elements"); } } } public void waitUntilDone() throws InterruptedException { join(); if (error != null) { throw new RuntimeException("Error in source thread", error); } } } }