MockLoginAzkabanServlet.java

/*
 * Copyright 2016 LinkedIn Corp.
 *
 * Licensed 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 azkaban.fixture;


import azkaban.server.AzkabanServer;
import azkaban.server.session.Session;
import azkaban.server.session.SessionCache;
import azkaban.user.UserManager;
import azkaban.utils.Props;
import azkaban.webapp.AzkabanWebServer;
import azkaban.webapp.servlet.LoginAbstractAzkabanServlet;
import azkaban.user.User;
import org.apache.velocity.app.VelocityEngine;
import org.mockito.Spy;
import org.mortbay.jetty.Server;

import javax.servlet.ServletConfig;
import javax.servlet.ServletException;
import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import java.io.IOException;
import java.util.UUID;

import static org.mockito.Mockito.*;

public class MockLoginAzkabanServlet extends LoginAbstractAzkabanServlet {

    private static final String SESSION_ID_NAME = "azkaban.browser.session.id";

    public static HttpServletRequest getRequestWithNoUpstream(String clientIp, String sessionId, String requestMethod){

        HttpServletRequest req = mock(HttpServletRequest.class);

        when(req.getRemoteAddr()).thenReturn(clientIp);
        when(req.getHeader("x-forwarded-for")).thenReturn(null);
        when(req.getMethod()).thenReturn(requestMethod);
        when(req.getContentType()).thenReturn("application/x-www-form-urlencoded");

        // Requires sessionId to be passed that is in the application's session cache
        when(req.getParameter("session.id")).thenReturn(sessionId);

        return req;
    }

    public static HttpServletRequest getRequestWithUpstream(String clientIp, String upstreamIp, String sessionId, String requestMethod){

        HttpServletRequest req = mock(HttpServletRequest.class);

        when(req.getRemoteAddr()).thenReturn("2.2.2.2:9999");
        when(req.getHeader("x-forwarded-for")).thenReturn(upstreamIp);
        when(req.getMethod()).thenReturn(requestMethod);
        when(req.getContentType()).thenReturn("application/x-www-form-urlencoded");

        // Requires sessionId to be passed that is in the application's session cache
        when(req.getParameter("session.id")).thenReturn(sessionId);

        return req;
    }

    public static HttpServletRequest getRequestWithMultipleUpstreams(String clientIp, String upstreamIp, String sessionId, String requestMethod){

        HttpServletRequest req = mock(HttpServletRequest.class);

        when(req.getRemoteAddr()).thenReturn("2.2.2.2:9999");
        when(req.getHeader("x-forwarded-for")).thenReturn(upstreamIp + ",1.1.1.1,3.3.3.3:33333");
        when(req.getMethod()).thenReturn(requestMethod);
        when(req.getContentType()).thenReturn("application/x-www-form-urlencoded");

        // Requires sessionId to be passed that is in the application's session cache
        when(req.getParameter("session.id")).thenReturn(sessionId);

        return req;
    }

    public static MockLoginAzkabanServlet getServletWithSession(String sessionId,
                                                                String username, String clientIp)
            throws Exception{

        MockLoginAzkabanServlet servlet = new MockLoginAzkabanServlet();

        Server server = mock(Server.class);
        Props props = new Props();
        UserManager userManager = mock(UserManager.class);

        // Need to mock and inject an application instance into the servlet
        AzkabanWebServer app = mock(AzkabanWebServer.class);

        MockLoginAzkabanServlet servletSpy = spy(servlet);

        when(servletSpy.getApplication()).thenReturn(app);

        // Create a concrete SessionCache so a session will get persisted
        // and can get looked up
        SessionCache cache = new SessionCache(props);
        when(app.getSessionCache()).thenReturn(cache);

        // Need a valid object here when processing a request
        when(app.getVelocityEngine()).thenReturn(mock(VelocityEngine.class));

        // Construct and store a session in the servlet
        azkaban.user.User user = mock(azkaban.user.User.class);
        when(user.getEmail()).thenReturn(username + "@mail.com");
        when(user.getUserId()).thenReturn(username);

        Session session = new Session(sessionId, user, clientIp);
        servletSpy.getApplication().getSessionCache().addSession(session);


        // Return the servletSpy since we replaced implementation for 'getApplication'
        return servletSpy;
    }

    @Override
    protected void handleGet(HttpServletRequest req, HttpServletResponse resp, Session session)
            throws ServletException, IOException {

        resp.getWriter().write("SUCCESS_MOCK_LOGIN_SERVLET");
    }

    @Override
    protected void handlePost(HttpServletRequest req, HttpServletResponse resp, Session session)
            throws ServletException, IOException {

        resp.getWriter().write("SUCCESS_MOCK_LOGIN_SERVLET");
    }
}