hbase-server/src/test/java/org/apache/hadoop/hbase/ConnectionRule.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one
   3  * or more contributor license agreements.  See the NOTICE file
   4  * distributed with this work for additional information
   5  * regarding copyright ownership.  The ASF licenses this file
   6  * to you under the Apache License, Version 2.0 (the
   7  * "License"); you may not use this file except in compliance
   8  * with the License.  You may obtain a copy of the License at
   9  *
  10  *     http://www.apache.org/licenses/LICENSE-2.0
  11  *
  12  * Unless required by applicable law or agreed to in writing, software
  13  * distributed under the License is distributed on an "AS IS" BASIS,
  14  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  15  * See the License for the specific language governing permissions and
  16  * limitations under the License.
  17  */
  18 package org.apache.hadoop.hbase;
  19
  20 import java.io.IOException;
  21 import java.util.concurrent.CompletableFuture;
  22 import java.util.function.Supplier;
  23 import org.apache.hadoop.hbase.client.AsyncConnection;
  24 import org.junit.ClassRule;
  25 import org.junit.Rule;
  26 import org.junit.rules.ExternalResource;
  27
  28 /**
  29  * A {@link Rule} that manages the lifecycle of an instance of {@link AsyncConnection}. Can be used
  30  * in either the {@link Rule} or {@link ClassRule} positions.
  31  * </p>
  32  * Use in combination with {@link MiniClusterRule}, for example:
  33  *
  34  * <pre>{@code
  35  *   public class TestMyClass {
  36  *     private static final MiniClusterRule miniClusterRule = MiniClusterRule.newBuilder().build();
  37  *     private static final ConnectionRule connectionRule =
  38  *       new ConnectionRule(miniClusterRule::createConnection);
  39  *
  40  *     @ClassRule
  41  *     public static final TestRule rule = RuleChain
  42  *       .outerRule(miniClusterRule)
  43  *       .around(connectionRule);
  44  *   }
  45  * }</pre>
  46  */
  47 public class ConnectionRule extends ExternalResource {
  48
  49   private final Supplier<CompletableFuture<AsyncConnection>> connectionSupplier;
  50   private AsyncConnection connection;
  51
  52   public ConnectionRule(final Supplier<CompletableFuture<AsyncConnection>> connectionSupplier) {
  53     this.connectionSupplier = connectionSupplier;
  54   }
  55
  56   public AsyncConnection getConnection() {
  57     return connection;
  58   }
  59
  60   @Override
  61   protected void before() throws Throwable {
  62     this.connection = connectionSupplier.get().join();
  63   }
  64
  65   @Override
  66   protected void after() {
  67     if (this.connection != null) {
  68       try {
  69         connection.close();
  70       } catch (IOException e) {
  71         throw new RuntimeException(e);
  72       }
  73     }
  74   }
  75 }