content: /*
   * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
   * (the "License"). You may not use this work except in compliance with the License, which is
   * available at www.apache.org/licenses/LICENSE-2.0
   *
   * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
   * either express or implied, as more fully set forth in the License.
   *
   * See the NOTICE file distributed with this work for information regarding copyright ownership.
   */
  
  package alluxio.testutils;
  
  import alluxio.AlluxioURI;
  import alluxio.AuthenticatedClientUserResource;
  import alluxio.conf.Configuration;
  import alluxio.conf.PropertyKey;
  import alluxio.grpc.DeletePOptions;
  import alluxio.master.LocalAlluxioCluster;
  import alluxio.master.file.FileSystemMaster;
  import alluxio.master.file.contexts.DeleteContext;
  import alluxio.master.file.contexts.ListStatusContext;
  import alluxio.metrics.MetricsSystem;
  import alluxio.security.authentication.AuthenticatedClientUser;
  import alluxio.underfs.UfsMode;
  import alluxio.util.SecurityUtils;
  import alluxio.wire.FileInfo;
  
  import org.junit.rules.TestRule;
  import org.junit.runner.Description;
  import org.junit.runners.model.Statement;
  
  import java.lang.annotation.Annotation;
  import java.lang.annotation.Retention;
  import java.lang.annotation.RetentionPolicy;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Map.Entry;
  import javax.annotation.concurrent.NotThreadSafe;
  
  /**
   * A JUnit Rule resource for automatically managing a local alluxio cluster for testing. To use it,
   * create an instance of the class under a {@literal @}Rule annotation, with the required
   * configuration parameters, and any necessary explicit {@link Configuration} settings. The
   * Alluxio cluster will be set up from scratch at the end of every method (or at the start of
   * every suite if {@literal @}ClassRule is used), and destroyed at the end. Below is an example
   * of declaring and using it.
   *
   * <pre>
   *   public class SomethingTest {
   *    {@literal @}Rule
   *    public LocalAlluxioClusterResource localAlluxioClusterResource =
   *      new LocalAlluxioClusterResource(WORKER_CAPACITY, BLOCK_SIZE);
   *
   *    {@literal @}Test
   *    public void testSomething() {
   *      localAlluxioClusterResource.get().getClient().create("/abced");
   *      ...
   *    }
   *
   *    {@literal @}Test
   *    {@literal @}LocalAlluxioClusterResource.Config(
   *        confParams = {CONF_KEY_1, CONF_VALUE_1, CONF_KEY_2, CONF_VALUE_2, ...})
   *    public void testSomethingWithDifferentConf() {
    *      localAlluxioClusterResource.get().getClient().create("/efghi");
   *      ...
   *    }
   *
   *    {@literal @}Test
   *    {@literal @}LocalAlluxioClusterResource.Config(startCluster = false)
   *    public void testSomethingWithClusterStartedManually() {
   *      localAlluxioClusterResource.start();
   *      localAlluxioClusterResource.get().getClient().create("/efghi");
   *      ...
   *    }
   *   }
   * </pre>
   */
  @NotThreadSafe
  public final class LocalAlluxioClusterResource implements TestRule {
  
    /** Number of Alluxio workers in the cluster. */
    private final int mNumWorkers;
  
    /** Weather to include the secondary master. */
    private final boolean mIncludeSecondary;
  
    /** Weather to include the proxy. */
    private final boolean mIncludeProxy;
  
    /**
     * If true (default), we start the cluster before running a test method. Otherwise, the method
     * must start the cluster explicitly.
     */
    private final boolean mStartCluster;
  
    /** Configuration values for the cluster. */
    private final Map<PropertyKey, Object> mConfiguration = new HashMap<>();
  
    /** The Alluxio cluster being managed. */
    private LocalAlluxioCluster mLocalAlluxioCluster = null;
  
    /** The name of the test/cluster. */
    private String mTestName = "test";
  
    /**
     * Creates a new instance.
     *
     * @param startCluster whether or not to start the cluster before the test method starts
     * @param numWorkers the number of Alluxio workers to launch
     * @param configuration configuration for configuring the cluster
     */
    private LocalAlluxioClusterResource(boolean startCluster, boolean includeSecondary,
        boolean includeProxy, int numWorkers, Map<PropertyKey, Object> configuration) {
      mStartCluster = startCluster;
      mIncludeSecondary = includeSecondary;
      mIncludeProxy = includeProxy;
      mNumWorkers = numWorkers;
      mConfiguration.putAll(configuration);
      MetricsSystem.resetCountersAndGauges();
    }
  
    /**
     * @return the {@link LocalAlluxioCluster} being managed
     */
    public LocalAlluxioCluster get() {
      return mLocalAlluxioCluster;
    }
  
    /**
     * Adds a property to the cluster resource. Unset a property by passing a null value.
     *
     * @param key property key
     * @param value property value
     * @return the cluster resource
     */
    public LocalAlluxioClusterResource setProperty(PropertyKey key, Object value) {
      if (value == null) {
        mConfiguration.remove(key);
      } else {
        mConfiguration.put(key, value);
      }
      return this;
    }
  
    /**
     * Explicitly starts the {@link LocalAlluxioCluster}.
     */
    public void start() throws Exception {
      AuthenticatedClientUser.remove();
      // Create a new cluster.
      mLocalAlluxioCluster = new LocalAlluxioCluster(mNumWorkers, mIncludeSecondary, mIncludeProxy);
      // Init configuration for integration test
      mLocalAlluxioCluster.initConfiguration(mTestName);
      // Overwrite the test configuration with test specific parameters
      for (Entry<PropertyKey, Object> entry : mConfiguration.entrySet()) {
        Configuration.set(entry.getKey(), entry.getValue());
      }
      Configuration.global().validate();
      // Start the cluster
      mLocalAlluxioCluster.start();
    }
  
    /**
     * Explicitly stops the {@link LocalAlluxioCluster}.
     */
    public void stop() throws Exception {
      mLocalAlluxioCluster.stop();
    }
  
    @Override
    public Statement apply(final Statement statement, Description description) {
      return new Statement() {
        @Override
        public void evaluate() throws Throwable {
          IntegrationTestUtils.reserveMasterPorts();
          mTestName = IntegrationTestUtils
              .getTestName(description.getTestClass().getSimpleName(), description.getMethodName());
          try {
            try {
              Annotation configAnnotation;
              configAnnotation = description.getAnnotation(ServerConfig.class);
              if (configAnnotation != null) {
                overrideConfiguration(((ServerConfig) configAnnotation).confParams());
              }
  
              boolean startCluster = mStartCluster;
              configAnnotation = description.getAnnotation(Config.class);
              if (configAnnotation != null) {
                Config config = (Config) configAnnotation;
                overrideConfiguration(config.confParams());
                // Override startCluster
                startCluster = config.startCluster();
              }
              if (startCluster) {
                start();
              }
            } catch (Exception e) {
              throw new RuntimeException(e);
            }
            try {
              statement.evaluate();
            } finally {
              stop();
            }
          } finally {
            IntegrationTestUtils.releaseMasterPorts();
          }
        }
      };
    }
  
    /**
     * @param config the array of strings for the configuration values to override
     */
    private void overrideConfiguration(String[] config) {
      // Override the configuration parameters with any configuration params
      for (int i = 0; i < config.length; i += 2) {
        PropertyKey key = PropertyKey.fromString(config[i]);
        mConfiguration.put(key, key.parseValue(config[i + 1]));
      }
    }
  
    /**
     * Returns a resource which will reset the cluster without restarting it. The rule will perform
     * operations on the running cluster to get back to a clean state. This is primarily useful for
     * when the {@link LocalAlluxioCluster} is a {@code @ClassRule}, so this reset rule can reset
     * the cluster between tests.
     *
     * @return a {@link TestRule} for resetting the cluster
     */
    public TestRule getResetResource() {
      return new ResetRule(this);
    }
  
    /**
     * Builder for a {@link LocalAlluxioClusterResource}.
     */
    public static class Builder {
      private boolean mStartCluster;
      private boolean mIncludeSecondary;
      private boolean mIncludeProxy;
      private int mNumWorkers;
      private Map<PropertyKey, Object> mConfiguration;
  
      /**
       * Constructs the builder with default values.
       */
      public Builder() {
        mStartCluster = true;
        mIncludeSecondary = false;
        mIncludeProxy = false;
        mNumWorkers = 1;
        mConfiguration = new HashMap<>();
      }
  
      /**
       * @param startCluster whether to start the cluster at the start of the test
       */
      public Builder setStartCluster(boolean startCluster) {
        mStartCluster = startCluster;
        return this;
      }
  
      /**
       * @param includeSecondary whether to include the secondary master
       */
      public Builder setIncludeSecondary(boolean includeSecondary) {
        mIncludeSecondary = includeSecondary;
        return this;
      }
  
      /**
       * @param includeProxy whether to include the proxy
       */
      public Builder setIncludeProxy(boolean includeProxy) {
        mIncludeProxy = includeProxy;
        return this;
      }
  
      /**
       * @param numWorkers the number of workers to run in the cluster
       */
      public Builder setNumWorkers(int numWorkers) {
        mNumWorkers = numWorkers;
        return this;
      }
  
      /**
       * @param key the property key to set for the cluster
       * @param value the value to set it to
       */
      public Builder setProperty(PropertyKey key, Object value) {
        mConfiguration.put(key, value);
        return this;
      }
  
      /**
       * @return a {@link LocalAlluxioClusterResource} for the current builder values
       */
      public LocalAlluxioClusterResource build() {
        return new LocalAlluxioClusterResource(mStartCluster, mIncludeSecondary, mIncludeProxy,
            mNumWorkers, mConfiguration);
      }
    }
  
    private final class ResetRule implements TestRule {
      private final LocalAlluxioClusterResource mCluster;
  
      ResetRule(LocalAlluxioClusterResource cluster) {
        mCluster = cluster;
      }
  
      @Override
      public Statement apply(Statement statement, Description description) {
        return new Statement() {
          @Override
          public void evaluate() throws Throwable {
            try {
              statement.evaluate();
            } finally {
              FileSystemMaster fsm =
                  mCluster.mLocalAlluxioCluster.getLocalAlluxioMaster().getMasterProcess()
                      .getMaster(FileSystemMaster.class);
  
              if (SecurityUtils.isAuthenticationEnabled(Configuration.global())) {
                // Reset the state as the root inode user (superuser).
                try (AuthenticatedClientUserResource r = new AuthenticatedClientUserResource(
                    fsm.getRootInodeOwner(), Configuration.global())) {
                  resetCluster(fsm);
                }
              } else {
                resetCluster(fsm);
              }
            }
          }
        };
      }
  
      private void resetCluster(FileSystemMaster fsm) throws Exception {
        if (!mCluster.get().getLocalAlluxioMaster().isServing()) {
          // Restart the masters if they are not serving.
          mCluster.mLocalAlluxioCluster.startMasters();
          // use the newly created/started file system master
          fsm = mCluster.mLocalAlluxioCluster.getLocalAlluxioMaster().getMasterProcess()
              .getMaster(FileSystemMaster.class);
        }
        if (!mCluster.get().isStartedWorkers()) {
          // Start the new workers
          mCluster.get().startWorkers();
        }
        fsm.updateUfsMode(new AlluxioURI(fsm.getUfsAddress()), UfsMode.READ_WRITE);
        for (FileInfo fileInfo : fsm
            .listStatus(new AlluxioURI("/"), ListStatusContext.defaults())) {
          fsm.delete(new AlluxioURI(fileInfo.getPath()), DeleteContext
              .create(DeletePOptions.newBuilder()
                  .setUnchecked(true)
                  .setRecursive(true)
                  .setDeleteMountPoint(true)));
        }
      }
    }
  
    /**
     * An annotation for test methods that can be used to override any of the defaults set at
     * construction time.
     */
    @Retention(RetentionPolicy.RUNTIME)
    public @interface Config {
      String[] confParams() default {};
      boolean startCluster() default true;
    }
  
    /**
     * Class-level annotations that can override the configuration for static class rules.
     */
    @Retention(RetentionPolicy.RUNTIME)
    public @interface ServerConfig {
      String[] confParams() default {};
    }
  }
  