ZeppelinClientTest.java

package org.apache.zeppelin.notebook.repo.zeppelinhub.websocket;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.mock.MockEchoWebsocketServer;
import org.apache.zeppelin.notebook.socket.Message;
import org.apache.zeppelin.notebook.socket.Message.OP;
import org.eclipse.jetty.websocket.api.Session;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.Maps;

public class ZeppelinClientTest {
  private Logger LOG = LoggerFactory.getLogger(ZeppelinClientTest.class);
  private final int zeppelinPort = 8080;
  private final String validWebsocketUrl = "ws://localhost:" + zeppelinPort + "/ws";
  private ExecutorService executor;
  private MockEchoWebsocketServer echoServer;

  @Before
  public void setUp() throws Exception {
    startWebsocketServer();
  }

  @After
  public void tearDown() throws Exception {
    //tear down routine
    echoServer.stop();
    executor.shutdown();
  }

  private void startWebsocketServer() throws InterruptedException {
    // mock zeppelin websocket server setup
    executor = Executors.newFixedThreadPool(1);
    echoServer = new MockEchoWebsocketServer(zeppelinPort);
    executor.submit(echoServer);
  }

  @Test
  public void zeppelinConnectionTest() {
    try {
      // Wait for websocket server to start
      Thread.sleep(2000);
    } catch (InterruptedException e) {
      LOG.warn("Cannot wait for websocket server to start, returning");
      return;
    }
    // Initialize and start Zeppelin client
    ZeppelinClient client = ZeppelinClient.initialize(validWebsocketUrl, "dummy token", null);
    client.start();
    LOG.info("Zeppelin websocket client started");

    // Connection to note AAAA
    Session connectionA = client.getZeppelinConnection("AAAA", "anonymous", "anonymous");
    assertNotNull(connectionA);
    assertTrue(connectionA.isOpen());

    assertEquals(client.countConnectedNotes(), 1);
    assertEquals(connectionA, client.getZeppelinConnection("AAAA", "anonymous", "anonymous"));

    // Connection to note BBBB
    Session connectionB = client.getZeppelinConnection("BBBB", "anonymous", "anonymous");
    assertNotNull(connectionB);
    assertTrue(connectionB.isOpen());

    assertEquals(client.countConnectedNotes(), 2);
    assertEquals(connectionB, client.getZeppelinConnection("BBBB", "anonymous", "anonymous"));

    // Remove connection to note AAAA
    client.removeNoteConnection("AAAA");
    assertEquals(client.countConnectedNotes(), 1);
    assertNotEquals(connectionA, client.getZeppelinConnection("AAAA", "anonymous", "anonymous"));
    assertEquals(client.countConnectedNotes(), 2);
    client.stop();
  }

  @Test
  public void zeppelinClientSingletonTest() {
    ZeppelinClient client1 = ZeppelinClient.getInstance();
    if (client1 == null) {
      client1 = ZeppelinClient.initialize(validWebsocketUrl, "TOKEN", null);
    }
    assertNotNull(client1);
    ZeppelinClient client2 = ZeppelinClient.getInstance();
    assertNotNull(client2);
    assertEquals(client1, client2);
  }

  @Test
  public void zeppelinMessageSerializationTest() {
    Message msg = new Message(OP.LIST_NOTES);
    msg.data = Maps.newHashMap();
    msg.data.put("key", "value");
    ZeppelinClient client = ZeppelinClient.initialize(validWebsocketUrl, "TOKEN", null);
    String serializedMsg = client.serialize(msg);
    Message deserializedMsg = client.deserialize(serializedMsg);
    assertEquals(msg.op, deserializedMsg.op);
    assertEquals(msg.data.get("key"), deserializedMsg.data.get("key"));
    
    String invalidMsg = "random text";
    deserializedMsg =client.deserialize(invalidMsg);
    assertNull(deserializedMsg);
  }

  @Test
  public void sendToZeppelinTest() {
    ZeppelinClient client = ZeppelinClient.initialize(validWebsocketUrl, "TOKEN", null);
    client.start();
    Message msg = new Message(OP.LIST_NOTES);
    msg.data = Maps.newHashMap();
    msg.data.put("key", "value");
    client.send(msg, "DDDD");
    client.removeNoteConnection("DDDD");
    client.stop();
  }
}