YarnSecurityManagerTest.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 gobblin.yarn;

import java.io.IOException;
import java.net.URL;
import java.util.Collection;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.RetryOneTime;
import org.apache.curator.test.TestingServer;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.helix.HelixManager;
import org.apache.helix.HelixManagerFactory;
import org.apache.helix.InstanceType;
import org.mockito.Mockito;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;

import com.google.common.base.Function;
import com.google.common.eventbus.EventBus;
import com.google.common.io.Closer;
import com.typesafe.config.Config;
import com.typesafe.config.ConfigFactory;
import com.typesafe.config.ConfigValueFactory;

import gobblin.cluster.GobblinClusterConfigurationKeys;
import gobblin.cluster.HelixUtils;
import gobblin.cluster.TestHelper;
import gobblin.testing.AssertWithBackoff;


/**
 * Unit tests for {@link YarnAppSecurityManager} and {@link YarnContainerSecurityManager}.
 *
 * <p>
 *   This class tests {@link YarnAppSecurityManager} and {@link YarnContainerSecurityManager} together
 *   as it is more convenient to test both here where all dependencies are setup between the two.
 * </p>
 *
 * <p>
 *   This class uses a {@link TestingServer} as an embedded ZooKeeper server for testing. The Curator
 *   framework is used to provide a ZooKeeper client. This class also uses a {@link HelixManager} as
 *   being required by {@link YarnAppSecurityManager}. The local file system as returned by
 *   {@link FileSystem#getLocal(Configuration)} is used for writing the testing delegation token, which
 *   is acquired by mocking the method {@link FileSystem#getDelegationToken(String)} on the local
 *   {@link FileSystem} instance.
 * </p>
 * @author Yinan Li
 */
@Test(groups = { "gobblin.yarn" })
public class YarnSecurityManagerTest {
  final Logger LOG = LoggerFactory.getLogger(YarnSecurityManagerTest.class);

  private CuratorFramework curatorFramework;

  private HelixManager helixManager;

  private Configuration configuration;
  private FileSystem localFs;
  private Path baseDir;
  private Path tokenFilePath;
  private Token<?> token;

  private YarnAppSecurityManager yarnAppSecurityManager;
  private YarnContainerSecurityManager yarnContainerSecurityManager;

  private final Closer closer = Closer.create();

  @BeforeClass
  public void setUp() throws Exception {
    // Use a random ZK port
    TestingServer testingZKServer = this.closer.register(new TestingServer(-1));
    LOG.info("Testing ZK Server listening on: " + testingZKServer.getConnectString());

    this.curatorFramework = this.closer.register(
        CuratorFrameworkFactory.newClient(testingZKServer.getConnectString(), new RetryOneTime(2000)));
    this.curatorFramework.start();

    URL url = YarnSecurityManagerTest.class.getClassLoader().getResource(
        YarnSecurityManagerTest.class.getSimpleName() + ".conf");
    Assert.assertNotNull(url, "Could not find resource " + url);

    Config config = ConfigFactory.parseURL(url)
        .withValue("gobblin.cluster.zk.connection.string",
                   ConfigValueFactory.fromAnyRef(testingZKServer.getConnectString()))
        .resolve();

    String zkConnectingString = config.getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
    String helixClusterName = config.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY);

    HelixUtils.createGobblinHelixCluster(zkConnectingString, helixClusterName);

    this.helixManager = HelixManagerFactory.getZKHelixManager(
        helixClusterName, TestHelper.TEST_HELIX_INSTANCE_NAME, InstanceType.SPECTATOR, zkConnectingString);
    this.helixManager.connect();

    this.configuration = new Configuration();
    this.localFs = Mockito.spy(FileSystem.getLocal(this.configuration));

    this.token = new Token<>();
    this.token.setKind(new Text("test"));
    this.token.setService(new Text("test"));
    Mockito.<Token<?>>when(this.localFs.getDelegationToken(UserGroupInformation.getLoginUser().getShortUserName()))
        .thenReturn(this.token);

    this.baseDir = new Path(YarnSecurityManagerTest.class.getSimpleName());
    this.tokenFilePath = new Path(this.baseDir, GobblinYarnConfigurationKeys.TOKEN_FILE_NAME);
    this.yarnAppSecurityManager =
        new YarnAppSecurityManager(config, this.helixManager, this.localFs, this.tokenFilePath);
    this.yarnContainerSecurityManager = new YarnContainerSecurityManager(config, this.localFs, new EventBus());
  }

  @Test
  public void testGetNewDelegationTokenForLoginUser() throws IOException {
    this.yarnAppSecurityManager.getNewDelegationTokenForLoginUser();
  }

  @Test(dependsOnMethods = "testGetNewDelegationTokenForLoginUser")
  public void testWriteDelegationTokenToFile() throws IOException {
    this.yarnAppSecurityManager.writeDelegationTokenToFile();
    Assert.assertTrue(this.localFs.exists(this.tokenFilePath));
    assertToken(YarnHelixUtils.readTokensFromFile(this.tokenFilePath, this.configuration));
  }


  static class GetControllerMessageNumFunc implements Function<Void, Integer> {
    private final CuratorFramework curatorFramework;
    private final String testName;

    public GetControllerMessageNumFunc(String testName, CuratorFramework curatorFramework) {
      this.curatorFramework = curatorFramework;
      this.testName = testName;
    }

    @Override
    public Integer apply(Void input) {
      try {
        return this.curatorFramework.getChildren().forPath(String.format("/%s/CONTROLLER/MESSAGES",
            this.testName)).size();
      } catch (Exception e) {
        throw new RuntimeException(e);
      }
    }

  }

  @Test
  public void testSendTokenFileUpdatedMessage() throws Exception {
    Logger log = LoggerFactory.getLogger("testSendTokenFileUpdatedMessage");
    this.yarnAppSecurityManager.sendTokenFileUpdatedMessage(InstanceType.CONTROLLER);
    Assert.assertEquals(this.curatorFramework.checkExists().forPath(
        String.format("/%s/CONTROLLER/MESSAGES", YarnSecurityManagerTest.class.getSimpleName())).getVersion(), 0);
    AssertWithBackoff.create().logger(log).timeoutMs(20000)
      .assertEquals(new GetControllerMessageNumFunc(YarnSecurityManagerTest.class.getSimpleName(),
          this.curatorFramework), 1, "1 controller message queued");
  }

  @Test(dependsOnMethods = "testWriteDelegationTokenToFile")
  public void testYarnContainerSecurityManager() throws IOException {
    Collection<Token<?>> tokens = this.yarnContainerSecurityManager.readDelegationTokens(this.tokenFilePath);
    assertToken(tokens);
    this.yarnContainerSecurityManager.addDelegationTokens(tokens);
    assertToken(UserGroupInformation.getCurrentUser().getTokens());
  }

  @AfterClass
  public void tearDown() throws IOException {
    try {
      if (this.helixManager.isConnected()) {
        this.helixManager.disconnect();
      }
      this.localFs.delete(this.baseDir, true);
    } catch (Throwable t) {
      throw this.closer.rethrow(t);
    } finally {
      this.closer.close();
    }
  }

  private void assertToken(Collection<Token<?>> tokens) {
    Assert.assertEquals(tokens.size(), 1);
    Token<?> token = tokens.iterator().next();
    Assert.assertEquals(token, this.token);
  }
}