package org.realityforge.replicant.server.ee.rest;

import javax.annotation.Nonnull;
import javax.persistence.EntityManager;
import javax.transaction.TransactionSynchronizationRegistry;
import javax.ws.rs.core.Response;
import org.realityforge.replicant.server.EntityMessageEndpoint;
import org.realityforge.replicant.server.transport.ReplicantSession;
import org.realityforge.replicant.server.transport.ReplicantSessionManager;
import org.testng.annotations.Test;
import static org.mockito.Mockito.*;
import static org.testng.Assert.*;

public class SessionRestServiceTest
{
  @Test
  public void poll_dataAvailable()
    throws Exception
  {
    final ReplicantSessionManager sessionManager = mock( ReplicantSessionManager.class );
    final AbstractSessionRestService resource = newResource( sessionManager );

    when( sessionManager.createSession() ).
      thenReturn( new ReplicantSession( null, "2222" ) );

    final Response token = resource.createSession();

    assertEquals( token.getEntity(), "2222" );
  }

  private AbstractSessionRestService newResource( final ReplicantSessionManager sessionManager )
    throws Exception
  {
    return new AbstractSessionRestService()
    {
      @Nonnull
      @Override
      protected EntityMessageEndpoint getEntityMessageEndpoint()
      {
        return mock( EntityMessageEndpoint.class );
      }

      @Nonnull
      @Override
      protected EntityManager getEntityManager()
      {
        return mock( EntityManager.class );
      }

      @Nonnull
      @Override
      protected ReplicantSessionManager getSessionManager()
      {
        return sessionManager;
      }

      @Nonnull
      @Override
      protected TransactionSynchronizationRegistry getRegistry()
      {
        return mock( TransactionSynchronizationRegistry.class );
      }
    };
  }
}

@@ -803,24 +803,97 @@
 Manager );%0A%0A
+    final ReplicantSession session = new ReplicantSession( null, null );%0A
     when( se
@@ -946,44 +946,15 @@
 rn( 
-new ReplicantSession( null, %222222%22 )
+session
  );%0A
@@ -1049,14 +1049,30 @@
 (), 
-%222222%22
+session.getSessionID()
  );%0A

Simplify test