

package com.netflix.lipstick.pigtolipstick;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Properties;

import org.testng.Assert;
import org.testng.annotations.Test;

import static org.mockito.Mockito.*;

import org.apache.pig.LipstickPigServer;
import org.apache.pig.impl.PigContext;
import com.netflix.lipstick.P2jPlanGenerator;
import com.netflix.lipstick.pigstatus.PigStatusClient;
import com.netflix.lipstick.model.P2jPlanPackage;
import com.netflix.lipstick.model.P2jPlan;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.plans.MROperPlan;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.MapReduceOper;
import org.apache.pig.backend.hadoop.executionengine.HExecutionEngine;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.PhysicalOperator;
import org.apache.pig.newplan.Operator;

public class BasicP2LClientTest {

    /* Create a Mock PigContext object with methods stubbed out to
       let the BasicP2LClient.createPlan() method complete */
    public PigContext makeMockContext() {
        HExecutionEngine exec_engine = mock(HExecutionEngine.class);
        when(exec_engine.getLogToPhyMap()).thenReturn(new HashMap<Operator, PhysicalOperator>());

        PigContext ctx = mock(PigContext.class);
        when(ctx.getExecutionEngine()).thenReturn(exec_engine);
        when(ctx.getProperties()).thenReturn(new Properties());
        return ctx;
    }

    /* Create Mock P2jPlanGenerator objects with methods stubbed out to
       let the BasicP2LClient.createPlan() method complete, and
       pass them to the client via the setter method. */
    public void setMockPlanGenerators(BasicP2LClient client) {
        P2jPlan pln = mock(P2jPlan.class);
        P2jPlanGenerator unopt_plan = mock(P2jPlanGenerator.class);
        when(unopt_plan.getP2jPlan()).thenReturn(pln);
        P2jPlanGenerator opt_plan = mock(P2jPlanGenerator.class);
        when(opt_plan.getP2jPlan()).thenReturn(pln);
        client.setPlanGenerators(unopt_plan, opt_plan);
    }

    public MROperPlan makeMockMROperPlan() {
        MROperPlan plan = mock(MROperPlan.class);
        when(plan.iterator()).thenReturn(new ArrayList<MapReduceOper>().iterator());
        return plan;
    }

    @Test
    public void testCreatePlanClientSaved() throws Exception {
        PigStatusClient status_client = mock(PigStatusClient.class);        
        BasicP2LClient client = new BasicP2LClient(status_client);
        
        LipstickPigServer server = mock(LipstickPigServer.class);
        PigContext ctx = makeMockContext();
        when(server.getPigContext()).thenReturn(ctx);
        client.setPigServer(server);

        setMockPlanGenerators(client);

        client.createPlan(makeMockMROperPlan());

        /* the client.savePlan() method should have been called */
        verify(status_client).savePlan(any(P2jPlanPackage.class));        
    }

    @Test
    public void testCreatePlanNoPigServerJustContext() throws Exception {
        PigStatusClient status_client = mock(PigStatusClient.class);        
        BasicP2LClient client = new BasicP2LClient(status_client);
        
        PigContext ctx = makeMockContext();
        client.setPigContext(ctx);

        setMockPlanGenerators(client);

        client.createPlan(makeMockMROperPlan());

        /* the client.savePlan() method should have been called */
        verify(status_client).savePlan(any(P2jPlanPackage.class));        
    }
}

@@ -904,551 +904,8 @@
  %7B%0A%0A
-    /* Create a Mock PigContext object with methods stubbed out to%0A       let the BasicP2LClient.createPlan() method complete */%0A    public PigContext makeMockContext() %7B%0A        HExecutionEngine exec_engine = mock(HExecutionEngine.class);%0A        when(exec_engine.getLogToPhyMap()).thenReturn(new HashMap%3COperator, PhysicalOperator%3E());%0A%0A        PigContext ctx = mock(PigContext.class);%0A        when(ctx.getExecutionEngine()).thenReturn(exec_engine);%0A        when(ctx.getProperties()).thenReturn(new Properties());%0A        return ctx;%0A    %7D%0A%0A
     
@@ -1704,32 +1704,777 @@
 rn plan;%0A    %7D%0A%0A
+    /* Create a Mock PigContext object with methods stubbed out to%0A        let the BasicP2LClient.createPlan() method complete */%0A    public void addMockPigContext(BasicP2LClient client) %7B%0A        HExecutionEngine exec_engine = mock(HExecutionEngine.class);%0A        when(exec_engine.getLogToPhyMap()).thenReturn(new HashMap%3COperator, PhysicalOperator%3E());%0A%0A        PigContext ctx = mock(PigContext.class);%0A        when(ctx.getExecutionEngine()).thenReturn(exec_engine);%0A        when(ctx.getProperties()).thenReturn(new Properties());%0A        client.setPigContext(ctx);%0A    %7D%0A%0A    public void addMockPigServer(BasicP2LClient client) %7B%0A        LipstickPigServer server = mock(LipstickPigServer.class);%0A        client.setPigServer(server);%0A    %7D%0A%0A%0A
     @Test%0A    pu
@@ -2670,32 +2670,24 @@
 us_client);%0A
-        
 %0A        Lip
@@ -2687,161 +2687,32 @@
     
-LipstickPigServer server = mock(LipstickPigServer.class);%0A        PigContext ctx = makeMockContext();%0A        when(server.getPigContext()).thenReturn(ctx
+addMockPigContext(client
 );%0A 
@@ -2714,34 +2714,31 @@
 t);%0A        
-client.set
+addMock
 PigServer(se
@@ -2739,18 +2739,17 @@
 ver(
-server
+client
 );%0A
-%0A
     
@@ -3223,50 +3223,26 @@
     
-%0A        PigContext ctx = make
+add
 Mock
+Pig
 Context(
 );%0A 
@@ -3241,47 +3241,17 @@
 ext(
-);%0A        client.setPigContext(ctx
+client
 );%0A
-%0A
     

update test to reflect change in function call semantics