class Assign implements Command {

    @Override
    public void run() throws Exception {
        String actionStr = ((String) args.getOptionSet().nonOptionArguments().get(1)).toUpperCase();
        AssignAction assignAction = AssignAction.valueOf(actionStr);
        assignAction.run(admin, args);
    }

}

class AssignTest  {

    @Test
    public void testExportImportWithBalancerOffOption() throws Exception {

        String expFileName = "export_test.exp";
        boolean balancerRunning = false;

        // invalid export
        try {
            balancerRunning = false;
            String[] argsParam = {"zookeeper", "export", expFileName};
            Args args = new ManagerArgs(argsParam);
            Assign command = new Assign(admin, args);

            balancerRunning = admin.setBalancerRunning(true, true);
            command.run();
            fail();
        } catch (IllegalStateException e) {
            if (!e.getMessage().contains(AssignAction.MESSAGE_TURN_BALANCER_OFF))
                throw e;
        } finally {
            if (balancerRunning)
                admin.setBalancerRunning(true, true);
        }

        // valid export
        try {
            balancerRunning = false;
            String[] argsParam = {"zookeeper", "export", expFileName, "--" + Args.OPTION_TURN_BALANCER_OFF,
                "--skip-export", "--force-proceed"};
            Args args = new ManagerArgs(argsParam);
            Assign command = new Assign(admin, args);

            balancerRunning = admin.setBalancerRunning(true, true);
            command.run();
        } finally {
            if (balancerRunning)
                admin.setBalancerRunning(true, true);
        }

        // invalid import
        try {
            balancerRunning = false;
            String[] argsParam = {"zookeeper", "import", expFileName};
            Args args = new ManagerArgs(argsParam);
            Assign command = new Assign(admin, args);

            balancerRunning = admin.setBalancerRunning(true, true);
            command.run();
            fail();
        } catch (IllegalStateException e) {
            if (!e.getMessage().contains(AssignAction.MESSAGE_TURN_BALANCER_OFF))
                throw e;
        } finally {
            if (balancerRunning)
                admin.setBalancerRunning(true, true);
        }

        // valid import
        try {
            balancerRunning = false;
            String[] argsParam = {"zookeeper", "import", expFileName, "--" + Args.OPTION_TURN_BALANCER_OFF,
                "--skip-export", "--force-proceed"};
            Args args = new ManagerArgs(argsParam);
            Assign command = new Assign(admin, args);

            balancerRunning = admin.setBalancerRunning(true, true);
            command.run();
        } finally {
            if (balancerRunning)
                admin.setBalancerRunning(true, true);
        }
    }

}