#
# Copyright 2011-2013 Amazon.com, Inc. or its affiliates.  All Rights Reserved.

require 'commands'
require 'test/unit'

module Commands

  class MockExecutor
    def exec(cmd)
    end
  end

  class MockEMRClient
    attr_accessor :state

    def initialize(config)
      @config = config
      @state = "RUNNING"
      @step_map =
      { "j-hive-installed" =>
        {"ExecutionStatusDetail"=>
          {"StartDateTime"=>1291074747.0, "EndDateTime"=>1291074776.0, "LastStateChangeReason"=>nil,
          "CreationDateTime"=>1291074521.0, "State"=>"COMPLETED"},
          "StepConfig"=>
            {"Name"=>"Setup Hive",
            "HadoopJarStep"=>
              {"Jar"=>"s3://us-east-1.elasticmapreduce/libs/script-runner/script-runner.jar",
              "Args"=>
                ["s3://us-east-1.elasticmapreduce/libs/hive/hive-script",
                "--base-path", "s3://us-east-1.elasticmapreduce/libs/hive/",
                "--install-hive"],
                "Properties"=>[],
                "MainClass"=>nil},
              "ActionOnFailure"=>"CONTINUE"}},
        "j-hive-0.5-installed" =>
        {"ExecutionStatusDetail"=>
          {"StartDateTime"=>1291074747.0, "EndDateTime"=>1291074776.0, "LastStateChangeReason"=>nil,
          "CreationDateTime"=>1291074521.0, "State"=>"COMPLETED"},
          "StepConfig"=>
            {"Name"=>"Setup Hive",
            "HadoopJarStep"=>
              {"Jar"=>"s3://us-east-1.elasticmapreduce/libs/script-runner/script-runner.jar",
              "Args"=>
                ["s3://us-east-1.elasticmapreduce/libs/hive/hive-script",
                "--base-path", "s3://us-east-1.elasticmapreduce/libs/hive/",
                "--install-hive","--hive-versions","0.5"],
                "Properties"=>[],
                "MainClass"=>nil},
              "ActionOnFailure"=>"CONTINUE"}},
        "j-hive-0.7-installed" =>
        {"ExecutionStatusDetail"=>
          {"StartDateTime"=>1291074747.0, "EndDateTime"=>1291074776.0, "LastStateChangeReason"=>nil,
          "CreationDateTime"=>1291074521.0, "State"=>"COMPLETED"},
          "StepConfig"=>
            {"Name"=>"Setup Hive",
            "HadoopJarStep"=>
              {"Jar"=>"s3://us-east-1.elasticmapreduce/libs/script-runner/script-runner.jar",
              "Args"=>
                ["s3://us-east-1.elasticmapreduce/libs/hive/hive-script",
                "--base-path", "s3://us-east-1.elasticmapreduce/libs/hive/",
                "--install-hive","--hive-versions","0.7"],
                "Properties"=>[],
                "MainClass"=>nil},
              "ActionOnFailure"=>"CONTINUE"}},
      }
    end

    def client
      return self
    end
    
    def describe_job_flows(args)
      steps = []
      if args["JobFlowIds"] != nil and @step_map.has_key? args["JobFlowIds"].first then
        steps = [@step_map[args["JobFlowIds"].first]]
      end
      return {
        "JobFlows" =>
        [
         {
           "LogUri" => "s3n://testing/", 
           "Name" => "Development Job Flow  (requires manual termination)", 
           "BootstrapActions" =>[], 
           "ExecutionStatusDetail" => {
             "EndDateTime" => 1286584312.0, 
             "CreationDateTime" => 1286584224.0, 
             "LastStateChangeReason" => "Terminated by user request", 
             "State" => @state, 
             "StartDateTime" => nil, 
             "ReadyDateTime" => nil
           }, 
           "Steps" => steps, 
           "JobFlowId" => "j-2HWO50OUKNMHG", 
           "Instances" => {
             "Ec2KeyName" => "richcole-test", 
             "InstanceCount" =>5, 
             "NormalizedInstanceHours" => nil, 
             "Placement" => {"AvailabilityZone" => "us-east-1d"}, 
             "KeepJobFlowAliveWhenNoSteps" => true, 
             "SlaveInstanceType" => "m2.xlarge", 
             "MasterInstanceType" => "m2.xlarge", 
             "MasterPublicDnsName" => nil, 
             "MasterInstanceId" => nil, 
             "InstanceGroups" => [{
               "SpotPrice" => nil,
               "EndDateTime" => nil,
               "Name" => "Task Instance Group",
               "InstanceRole" => "TASK",
               "CreationDateTime" => 1286862675.0,
               "LaunchGroup" => nil,
               "LastStateChangeReason" => "",
               "InstanceGroupId" => "ig-D2NC23WFSOOU",
               "State" => "RUNNING",
               "Market" => "ON_DEMAND",
               "InstanceType" => "c1.medium",
               "StartDateTime" => 1286862907.0,
               "InstanceRunningCount" => 2,
               "ReadyDateTime" => 1286862907.0,
               "InstanceRequestCount" => 2
               },
               {
               "SpotPrice" => nil,
               "EndDateTime" => nil,
               "Name" => "Master Instance Group",
               "InstanceRole" => "MASTER",
               "CreationDateTime" => 1286862675.0,
               "LaunchGroup" => nil,
               "LastStateChangeReason" => "",
               "InstanceGroupId" => "ig-1BFN7TCX7YE5Y",
               "State" => "RUNNING",
               "Market" => "ON_DEMAND",
               "InstanceType" => "m1.small",
               "StartDateTime" => 1286862866.0,
               "InstanceRunningCount" => 1,
               "ReadyDateTime" => 1286862906.0,
               "InstanceRequestCount" => 1
               },
               {
               "SpotPrice" => nil,
               "EndDateTime" => nil,
               "Name" => "Core Instance Group",
               "InstanceRole" => "CORE",
               "CreationDateTime" => 1286862675.0,
               "LaunchGroup" => nil,
               "LastStateChangeReason" => "Expanding cluster",
               "InstanceGroupId" => "ig-2EUIGTIPDLTXW",
               "State" => "RESIZING",
               "Market" => "ON_DEMAND",
               "InstanceType" => "m1.large",
               "StartDateTime" => 1286862907.0,
               "InstanceRunningCount" => 1,
               "ReadyDateTime" => 1286862907.0,
               "InstanceRequestCount" => 3
               }]
           },
           "HadoopVersion" => "0.20"
         } 
        ]
      }
    end

    def run_job_flow(opts)
      return { "JobFlowId" => "j-ABABABABA" }
    end

    def add_job_flow_steps(opts)
      return nil
    end

    def terminate_job_flows(opts)
      return nil
    end

    def modify_instance_groups(opts)
      return nil
    end

    def add_instance_groups(opts)
      return nil
    end

  end

  class MockLogger
    def puts(msg)
    end

    def trace(msg)
    end
  end

  class CommandTest < Test::Unit::TestCase

    def setup
      @client_class = MockEMRClient #FIXME: make this return always the same object
      @logger       = MockLogger.new
      @executor     = MockExecutor.new
    end

    def create_and_execute_commands(args)
      return ::Commands.create_and_execute_commands(args.split(/\s+/), @client_class, @logger, @executor, false)
    end

    def test_modify_instance_group_command
      args = "-c tests/credentials.json --modify-instance-group core --instance-count 10 --jobflow j-ABABABA"
      @commands = create_and_execute_commands(args)
      assert_equal(1, @commands.size)
      c = @commands.last
      assert(c.is_a?(ModifyInstanceGroupCommand))
      assert_equal(10, c.instance_count)
      assert_not_nil(c.instance_group_id)
      assert_equal(nil, c.instance_type)
      assert_equal("CORE", c.instance_role)
    end

    def test_one
      args = "-c tests/credentials.json --create --alive --num-instances 10 " +
        "--slave-instance-type m1.small --master-instance-type m1.large"
      @commands = create_and_execute_commands(args)
    end

    def test_two
      args = "-c tests/credentials.json --create --alive --num-instances 10 " +
        "--slave-instance-type m1.small --master-instance-type m1.large " + 
        "--instance-group TASK --instance-type m1.small --instance-count 10 " + 
        "--bootstrap-action s3://elasticmapreduce/scripts/configure-hadoop " + 
        "--arg s3://mybucket/config/custom-site-config.xml "
      @commands = create_and_execute_commands(args)
      assert_equal(1, @commands.commands.size)
    end

    def test_three
      args = "-c tests/credentials.json --create --alive --num-instances 10 " + 
        "--slave-instance-type m1.small --master-instance-type m1.large " +
        "--instance-group TASK --instance-type m1.small --instance-count 10 " + 
        "--bootstrap-action s3://elasticmapreduce/scripts/configure-hadoop " + 
        "--arg s3://mybucket/config/custom-site-config.xml " +
        "--pig-script s3://elasticmapreduce/samples/sample.pig " + 
        "--pig-interactive"
      @commands = create_and_execute_commands(args)
      assert_equal(1, @commands.commands.size)
      cmd1 = @commands.commands.first
      assert_equal(2, cmd1.step_commands.size)
      assert_equal(PigInteractiveCommand, cmd1.step_commands[0].class)
    end

    def test_four
      args = "-a ACCESS_ID -p SECRET_KEY --create --alive " + 
        "--hive-script s3://maps.google.com --enable-debugging " + 
        "--log-uri s3://somewhere.com/logs/"
      @commands = create_and_execute_commands(args)
      assert_equal(1, @commands.commands.size)
      assert_equal(3, @commands.commands[0].step_commands.size)
      steps = @commands.commands[0].step_commands
      assert_equal(EnableDebuggingCommand, steps[0].class) 
      assert_equal(HiveInteractiveCommand, steps[1].class) 
      assert_equal(HiveScriptCommand, steps[2].class) 
    end
    
    def hadoop_jar_step_args(command)
      command.steps[0]["HadoopJarStep"]["Args"]
    end
    
    def pig_script_command_default_args
      [
        "s3://us-east-1.elasticmapreduce/libs/pig/pig-script",
        "--base-path",
        "s3://us-east-1.elasticmapreduce/libs/pig/",
        "--run-pig-script",
        "--args",
        "-f"
      ]
    end
    
    def hive_script_command_default_args
      [
        "s3://us-east-1.elasticmapreduce/libs/hive/hive-script",
        "--base-path",
        "s3://us-east-1.elasticmapreduce/libs/hive/",
        "--run-hive-script",
        "--args",
        "-f"
      ]
    end

    def test_pig_arg
      args1 = "-a ACCESS_ID -p SECRET_KEY --create --alive " +
        "--pig-script --args s3://maps.google.com --args -p,INPUT=s3://maps.google.com/test.pig"
      args2 = "-a ACCESS_ID -p SECRET_KEY --create --alive " +
        "--pig-script s3://maps.google.com --args -p,INPUT=s3://maps.google.com/test.pig"
      @commands = create_and_execute_commands(args1)
      steps = @commands.commands.first.step_commands
      pig_command = steps[1]      
      assert_equal(pig_command.arg, nil)
      assert_equal(pig_command.args, ["s3://maps.google.com", "-p", "INPUT=s3://maps.google.com/test.pig"])
      assert_equal(pig_command.steps.size, 1)
      args = pig_script_command_default_args
      args.insert(3, "--pig-versions")
      args.insert(4, "latest")
      args << "s3://maps.google.com" << "-p" << "INPUT=s3://maps.google.com/test.pig"
      assert_equal(hadoop_jar_step_args(pig_command), args)
      
      @commands = create_and_execute_commands(args2)
      steps = @commands.commands.first.step_commands
      pig_command = steps[1]     
      assert_equal(pig_command.arg, "s3://maps.google.com")
      assert_equal(pig_command.args, ["-p", "INPUT=s3://maps.google.com/test.pig"])
      assert_equal(pig_command.steps.size, 1)
      assert_equal(hadoop_jar_step_args(pig_command), args)
    end
    
    def test_hive
      args = "-a ACCESS_ID -p SECRET_KEY --create --alive " + 
        "--hive-script --args s3://maps.google.com "
      @commands = create_and_execute_commands(args)
      create_command = @commands.commands.first
      steps = create_command.step_commands
      install_hive_command = steps[0]
      hive_command = steps[1]
      
      assert_equal(1, @commands.commands.size)
      assert_equal(2, steps.size)
      assert_equal(HiveScriptCommand, hive_command.class) 
      assert_equal(hive_command.arg, nil)
      assert_equal(hive_command.args, ["s3://maps.google.com"])
      assert_equal(hive_command.steps.size, 1)
      args = hive_script_command_default_args
      args.insert(3, "--hive-versions")
      args.insert(4, "latest")
      args << "s3://maps.google.com" 
      assert_equal(hadoop_jar_step_args(hive_command), args)
    end
    
    def test_install_hive_version
      args = "-a ACCESS_ID -p SECRET_KEY --create --alive " + 
        "--hive-script --args s3://maps.google.com --hive-versions 0.5"
      @commands = create_and_execute_commands(args)
      create_command = @commands.commands.first
      steps = create_command.step_commands
      install_hive_command = steps[0]
      hive_command = steps[1]
      assert_equal(1, @commands.commands.size)
      assert_equal(2, steps.size)
      assert_equal(HiveScriptCommand, hive_command.class) 
      assert_equal(hive_command.arg, nil)
      assert_equal(hive_command.args, ["s3://maps.google.com"])
      assert_equal(hive_command.steps.size, 1)
      args = hive_script_command_default_args
      args.insert(3, "--hive-versions")
      args.insert(4, "0.5")
      args  << "s3://maps.google.com"
      assert_equal(hadoop_jar_step_args(hive_command), args)
    end

    def test_run_hive_script_same_version
      args = "-a ACCESS_ID -p SECRET_KEY --jobflow j-hive-0.5-installed " + 
        "--hive-script --args s3://maps.google.com --hive-versions 0.5"
      @commands = create_and_execute_commands(args)
      create_command = @commands.commands.first
      steps = create_command.step_commands
      install_hive_command = steps[0]
      hive_command = steps[0]

      assert_equal(1, @commands.commands.size)
      assert_equal(1, steps.size)
      assert_equal(HiveScriptCommand, hive_command.class) 
      assert_equal(hive_command.arg, nil)
      assert_equal(hive_command.args, ["s3://maps.google.com"])
      assert_equal(hive_command.steps.size, 1)
      args = hive_script_command_default_args
      args.insert(3, "--hive-versions")
      args.insert(4, "0.5")
      args << "s3://maps.google.com"
      assert_equal(hadoop_jar_step_args(hive_command), args)
    end

    def test_run_hive_script_different_version
      args = "-a ACCESS_ID -p SECRET_KEY --jobflow j-hive-0.5-installed " + 
        "--hive-script --args s3://maps.google.com --hive-versions 0.7"
      @commands = create_and_execute_commands(args)
      create_command = @commands.commands.first
      steps = create_command.step_commands
      install_hive_command = steps[0]
      hive_command = steps[1]

      assert_equal(1, @commands.commands.size)
      assert_equal(2, steps.size)
      assert_equal(HiveScriptCommand, hive_command.class) 
      assert_equal(hive_command.arg, nil)
      assert_equal(hive_command.args, ["s3://maps.google.com"])
      assert_equal(hive_command.steps.size, 1)
      args = hive_script_command_default_args
      args.insert(3, "--hive-versions")
      args.insert(4, "0.7")
      args << "s3://maps.google.com"
      assert_equal(hadoop_jar_step_args(hive_command), args)
    end

    def test_hive_script_step_action_propogation
      args = "-a ACCESS_ID -p SECRET_KEY --jobflow j-hive-0.5-installed " + 
        "--hive-script --args s3://maps.google.com --hive-versions 0.7 " +
        "--step-action CONTINUE"
      @commands = create_and_execute_commands(args)
      create_command = @commands.commands.first
      steps = create_command.step_commands
      install_hive_command = steps[0]
      hive_command = steps[0]
      assert_equal(1, @commands.commands.size)
      assert_equal(2, steps.size)
      assert_equal(HiveInteractiveCommand, hive_command.class) 
      assert_equal(hive_command.arg, nil)
      assert_equal(hive_command.args, [])
      assert_equal(hive_command.steps.size, 1)
      assert_equal(hive_command.steps.first["ActionOnFailure"], "CONTINUE")
    end

    def test_hive_no_create
      args = "-a ACCESS_ID -p SECRET_KEY --jobflow j-ABABABABA " + 
        "--hive-script --args s3://maps.google.com "
      @commands = create_and_execute_commands(args)
      create_command = @commands.commands.first
      steps = create_command.step_commands
      install_hive_command = steps[0]
      hive_command = steps[1]

      assert_equal(1, @commands.commands.size)
      assert_equal(2, steps.size)
      assert_equal(HiveScriptCommand, hive_command.class) 
      assert_equal(hive_command.arg, nil)
      assert_equal(hive_command.args, ["s3://maps.google.com"])
      assert_equal(hive_command.steps.size, 1)
      args = hive_script_command_default_args
      args.insert(3, "--hive-versions")
      args.insert(4, "latest")
      args << "s3://maps.google.com" 
      assert_equal(hadoop_jar_step_args(hive_command), args)
    end
    
    def test_hive_no_create2
      args = "-a ACCESS_ID -p SECRET_KEY --jobflow j-ABABABABA " + 
        "--hive-script s3://maps.google.com --args -d,options "
      @commands = create_and_execute_commands(args)
      create_command = @commands.commands.first
      steps = create_command.step_commands
      install_hive_command = steps[0]
      hive_command = steps[1]

      assert_equal(1, @commands.commands.size)
      assert_equal(2, steps.size)
      assert_equal(HiveScriptCommand, hive_command.class) 
      assert_equal(hive_command.arg, "s3://maps.google.com")
      assert_equal(hive_command.args, ["-d", "options"])
      assert_equal(hive_command.steps.size, 1)
      args = hive_script_command_default_args
      args.insert(3, "--hive-versions")
      args.insert(4, "latest")
      args << "s3://maps.google.com" << "-d" << "options"
      assert_equal(hadoop_jar_step_args(hive_command), args)
    end   

    def test_five
      args = "-a ACCESS_ID -p SECRET_KEY -j j-ABABABAABA --hive-script " + 
        "s3://maps.google.com --enable-debugging --log-uri s3://somewhere.com/logs/"
      @commands = create_and_execute_commands(args)
      assert_equal(1, @commands.commands.size)
      assert_equal(3, @commands.commands[0].step_commands.size)
      steps = @commands.commands[0].step_commands
      assert_equal(EnableDebuggingCommand, steps[0].class) 
      assert_equal(HiveInteractiveCommand, steps[1].class) 
      assert_equal(HiveScriptCommand, steps[2].class) 
    end

    def test_six
      args = "-a ACCESS_ID -p SECRET_KEY --list --active"
      @commands = create_and_execute_commands(args)
    end

    def test_seven
      args = "-a ACCESS_ID -p SECRET_KEY --list --active --terminate"
      @commands = create_and_execute_commands(args)
    end

    def test_eight
      args = "-a ACCESS_ID -p SECRET_KEY --terminate -j j-ABABABABA"
      @commands = create_and_execute_commands(args)
    end

    def test_create_one
      args = "-a ACCESS_ID -p SECRET_KEY --create --alive --name TestFlow"
      @commands = create_and_execute_commands(args)
    end

    def test_ssh_no_jobflow
      args = "-a ACCESS_ID -p SECRET_KEY --ssh"
      assert_raise RuntimeError do
        @commands = create_and_execute_commands(args)
      end
    end

    def test_ssh_too_many_jobflows
      args = "-a ACCESS_ID -p SECRET_KEY -j j-ABABABA j-ABABABA --ssh"
      assert_raise RuntimeError do
        @commands = create_and_execute_commands(args)
      end
    end

    def test_jar_with_mainclass
      args = "-a ACCESS_ID -p SECRET_KEY -j j-3TRNB9E4GU2NI 
        --jar s3://my-example-bucket/wordcount.jar 
        --main-class org.myorg.WordCount 
        --arg s3://elasticmapreduce/samples/wordcount/input/ 
        --arg hdfs:///wordcount/output/1 
      "

      @commands = create_and_execute_commands(args)
      assert_equal(1, @commands.commands.size)
      c = @commands.commands.first
      s = c.step_commands.first
      st = s.steps
      assert_equal([{"HadoopJarStep"=>{"Jar"=>"s3://my-example-bucket/wordcount.jar",
        "Args"=>["s3://elasticmapreduce/samples/wordcount/input/", "hdfs:///wordcount/output/1"],
        "MainClass"=>"org.myorg.WordCount"}, "ActionOnFailure"=>"CANCEL_AND_WAIT", "Name"=>"Example Jar Step"}], st)
    end

    def test_ssh
      args = "-a ACCESS_ID -p SECRET_KEY --key-pair-file test.pem -j j-ABABABA --ssh"
      @commands = create_and_execute_commands(args)
    end

    def test_unarrest
      args = "-a ACCESS_ID -p SECRET_KEY --unarrest-instance-group core -j j-ABABABA"
      @commands = create_and_execute_commands(args)
    end

    def test_late_name
      args = "-a ACCESS_ID -p SECRET_KEY --create --alive --enable-debugging --hive-interactive --name MyHiveJobFlow --log-uri=s3://haijun-test/logs"
      @commands = create_and_execute_commands(args)
      assert_equal(1, @commands.commands.size)
      assert_equal("MyHiveJobFlow", @commands.commands.first.jobflow_name)
    end

    def test_ic_it
      args = "-a ACCESS_ID -p SECRET_KEY --create --alive --enable-debugging --hive-interactive --instance-count 5 --instance-type m1.small --name MyHiveJobFlow --log-uri=s3://haijun-test/logs"
      @commands = create_and_execute_commands(args)
      assert_equal(1, @commands.commands.size)
      cc = @commands.commands.first
      assert_equal("MyHiveJobFlow", cc.jobflow_name)
      assert_equal(5, cc.instance_count)
      assert_equal("m1.small", cc.instance_type)
    end

    def test_json
      args = "-a ACCESS_ID -p SECRET_KEY --jobflow j-ABABABABA --json tests/example.json --param <bucket>=mybucket --param mybucket=yourbucket"
      @commands = create_and_execute_commands(args)
      assert_equal(1, @commands.commands.size)
      cc = @commands.commands.first.step_commands.first
      assert_equal("tests/example.json", cc.arg)
      assert_equal({:key => "<bucket>", :value => "mybucket"}, cc.variables[0])
      assert_equal({:key => "mybucket", :value => "yourbucket"}, cc.variables[1])

      st = cc.steps.first

      expected_step = {
        "HadoopJarStep" => {
          "Jar"  => "/home/hadoop/contrib/streaming/hadoop-0.18-streaming.jar", 
          "Args" => ["-input", "s3n://elasticmapreduce/samples/wordcount/input", 
                     "-output", "s3n://yourbucket/result", "-mapper", "s3://yourbucket/lib/mapper"]
        }, 
        "ActionOnFailure" => "CONTINUE", 
        "Name" => "Example Step"
      }

      assert_equal(expected_step, st)
    end

    def test_ic_it2
      args = "-a ACCESS_ID -p SECRET_KEY --jobflow j-ABABABAB --resize-jobflow --add-instance-group task --instance-type m1.large --instance-count 5"
      @commands = create_and_execute_commands(args)
      assert_equal(2, @commands.commands.size)
      cc = @commands.commands[1]
      assert_equal(5, cc.instance_count)
      assert_equal("m1.large", cc.instance_type)
      assert_equal("TASK", cc.instance_role)
    end

    def test_command_option_mismatch
      args = "-c tests/credentials.json --instance-group core --instance-count 10"
      assert_raise RuntimeError do
        @commands = create_and_execute_commands(args)
      end
    end
    
    def test_region_from_az
      @commands = create_and_execute_commands("-c tests/credentials.json")
      eip_command = EipCommand.new("eip-command", "eip-command", "arg", @commands)
      assert_equal('https://ec2.us-east-1.amazonaws.com', eip_command.ec2_endpoint_from_az('us-east-1a'))
      assert_equal('https://ec2.ap-northeast-1.amazonaws.com', eip_command.ec2_endpoint_from_az('ec2.ap-northeast-1b'))
      assert_equal('https://ec2.us-west-1.amazonaws.com', eip_command.ec2_endpoint_from_az('ec2.us-west-1b'))
      assert_equal('https://ec2.us-west-2.amazonaws.com', eip_command.ec2_endpoint_from_az('ec2.us-west-2b'))
      assert_equal('https://ec2.sa-east-1.amazonaws.com', eip_command.ec2_endpoint_from_az('ec2.sa-east-1b'))
    end

    def test_hbase_instance_types
      assert_raise RuntimeError do
        @commands = create_and_execute_commands("-c tests/credentials.json --create --hbase")
      end
    end

  end
end
