package com.google.sps.servlets;

import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.After;

import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import static com.google.sps.tool.Tool.compareJson;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.*;

import java.io.PrintWriter;
import java.io.StringWriter;
import java.io.IOException;

import javax.servlet.annotation.WebServlet;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import com.google.appengine.tools.development.testing.LocalServiceTestHelper;
import com.google.appengine.tools.development.testing.LocalUserServiceTestConfig;
import com.google.appengine.tools.development.testing.LocalDatastoreServiceTestConfig;
import com.google.appengine.api.users.UserService;
import com.google.appengine.api.users.UserServiceFactory;
import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.DatastoreServiceFactory;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.KeyFactory;
import com.google.sps.data.User;
import com.google.gson.Gson;
import java.util.Map;
import java.util.HashMap;
import com.google.common.collect.ImmutableMap;


@RunWith(JUnit4.class)
public final class LoginServletTest {

  private static final String USER_ID = "testID";
  
  private final LocalServiceTestHelper helper =
    new LocalServiceTestHelper(
      new LocalDatastoreServiceTestConfig()
        .setDefaultHighRepJobPolicyUnappliedJobPercentage(0),
      new LocalUserServiceTestConfig())
      .setEnvIsAdmin(true).setEnvIsLoggedIn(true)
      .setEnvEmail("test@gmail.com").setEnvAuthDomain("gmail.com")
      .setEnvAttributes(
        new HashMap(
          ImmutableMap.of(
            "com.google.appengine.api.users.UserService.user_id_key", USER_ID)));

  private static final User USER_A = new User(USER_ID, "test@gmail.com");
  private static final User USER_B = new User("null", "null");
  private static final String LOGOUTURL = "/_ah/logout?continue\u003d%2F";
  private static final String LOGINURL = "/_ah/login?continue\u003d%2F";
  private HttpServletRequest mockRequest;
  private HttpServletResponse mockResponse;
  private StringWriter responseWriter;
  private LoginServlet servlet;
  private DatastoreService datastore;


  @Before
  public void setUp() throws Exception {
    helper.setUp();
    servlet = new LoginServlet();
    mockRequest = mock(HttpServletRequest.class);
    mockResponse = mock(HttpServletResponse.class);
    
    // Set up a fake HTTP response 
    responseWriter = new StringWriter();
    when(mockResponse.getWriter()).thenReturn(new PrintWriter(responseWriter));

    // Initialize datastore
    datastore = DatastoreServiceFactory.getDatastoreService();
  }

  @After
  public void tearDown() throws Exception {
    helper.setEnvIsLoggedIn(true);
    helper.tearDown();
  }
  
  @Test
  public void loginWhenUserInDatastoreAlready() throws Exception {
    
    storeUsertoDatastore(datastore, USER_A);

    servlet.doGet(mockRequest, mockResponse);
    String response = responseWriter.toString();
    
    Key key = USER_A.createEntity().getKey(); 
    USER_A.setUserKey(KeyFactory.keyToString(key));
    String expectedResponse = new Gson().toJson(getExpectedJsonInfo(USER_A, true, LOGOUTURL, "null"));

    assertTrue(compareJson(response, expectedResponse));
  }

  @Test
  public void userLoggedInButNotInDatastore() throws Exception {

    servlet.doGet(mockRequest, mockResponse);
    String response = responseWriter.toString();

    Key key = USER_A.createEntity().getKey(); 
    USER_A.setUserKey(KeyFactory.keyToString(key));
    String expectedResponse = new Gson().toJson(getExpectedJsonInfo(USER_A, true, LOGOUTURL, "null"));

    assertTrue(compareJson(response, expectedResponse));
  }

  @Test
  public void userNotLoggedIn() throws Exception {

    helper.setEnvIsLoggedIn(false);
    servlet.doGet(mockRequest, mockResponse);
    String response = responseWriter.toString();
    String expectedResponse = new Gson().toJson(getExpectedJsonInfo(USER_B, false, "null", LOGINURL));

    assertTrue(compareJson(response, expectedResponse));
  }
  
  public Map<String, Object> getExpectedJsonInfo(User user, boolean tabStatus, String logoutUrl, String loginUrl) {

    Map<String, Object> expectedJsonInfo = new HashMap<>();
    expectedJsonInfo.put("userInfo", user);
    expectedJsonInfo.put("showTabStatus", tabStatus);
    expectedJsonInfo.put("logoutUrl", logoutUrl);
    expectedJsonInfo.put("loginUrl", loginUrl);

    return expectedJsonInfo;
  }

  public void storeUsertoDatastore(DatastoreService datastore, User user) {
    datastore.put(user.createEntity());
  }
}
@@ -2017,22 +2017,30 @@
 al User 
+LOGGED_IN_
 USER
-_A
  = new U
@@ -2099,22 +2099,31 @@
 al User 
+LOGGED_OUT_
 USER
-_B
  = new U
@@ -3098,21 +3098,23 @@
 oid 
-loginWhenUser
+userLoggedInAnd
 InDa
@@ -3124,15 +3124,8 @@
 tore
-Already
 () t
@@ -3142,21 +3142,16 @@
 ption %7B%0A
-    %0A
     stor
@@ -3178,22 +3178,30 @@
 astore, 
+LOGGED_IN_
 USER
-_A
 );%0A%0A    
@@ -3302,38 +3302,46 @@
  %0A    Key key = 
+LOGGED_IN_
 USER
-_A
 .createEntity().
@@ -3347,38 +3347,46 @@
 .getKey(); %0A    
+LOGGED_IN_
 USER
-_A
 .setUserKey(KeyF
@@ -3483,33 +3483,95 @@
 nfo(
-USER_A, true, LOGOUTURL, 
+%0A       /*user*/LOGGED_IN_USER, /*tabStatus*/true, /*logoutUrl*/LOGOUTURL, /*loginUrl*/
 %22nul
@@ -3708,25 +3708,24 @@
 Exception %7B%0A
-%0A
     servlet.
@@ -3822,22 +3822,30 @@
 y key = 
+LOGGED_IN_
 USER
-_A
 .createE
@@ -3867,22 +3867,30 @@
 ); %0A    
+LOGGED_IN_
 USER
-_A
 .setUser
@@ -3995,33 +3995,96 @@
 nfo(
-USER_A, true, LOGOUTURL, 
+%0A        /*user*/LOGGED_IN_USER, /*tabStatus*/true, /*logoutUrl*/LOGOUTURL, /*loginUrl*/
 %22nul
@@ -4211,17 +4211,16 @@
 ption %7B%0A
-%0A
     help
@@ -4414,31 +4414,95 @@
 nfo(
-USER_B, false, %22null%22, 
+%0A        /*user*/LOGGED_OUT_USER, /*tabStatus*/false, /*logoutUrl*/%22null%22, /*loginUrl*/
 LOGI
@@ -4690,17 +4690,16 @@
 nUrl) %7B%0A
-%0A
     Map%3C

Add comments in getExpectedJsonInfo for readability