hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestTableSpaceQuotaViolationNotifier.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.quotas;
  19
  20 import static org.mockito.Matchers.argThat;
  21 import static org.mockito.Mockito.mock;
  22 import static org.mockito.Mockito.verify;
  23 import static org.mockito.Mockito.when;
  24
  25 import java.util.Arrays;
  26 import java.util.List;
  27 import java.util.Map.Entry;
  28 import java.util.NavigableMap;
  29 import java.util.Objects;
  30 import org.apache.hadoop.hbase.Cell;
  31 import org.apache.hadoop.hbase.HBaseClassTestRule;
  32 import org.apache.hadoop.hbase.TableName;
  33 import org.apache.hadoop.hbase.client.Connection;
  34 import org.apache.hadoop.hbase.client.Mutation;
  35 import org.apache.hadoop.hbase.client.Put;
  36 import org.apache.hadoop.hbase.client.Table;
  37 import org.apache.hadoop.hbase.quotas.SpaceQuotaSnapshot.SpaceQuotaStatus;
  38 import org.apache.hadoop.hbase.testclassification.SmallTests;
  39 import org.apache.hadoop.hbase.util.Bytes;
  40 import org.junit.Before;
  41 import org.junit.ClassRule;
  42 import org.junit.Test;
  43 import org.junit.experimental.categories.Category;
  44 import org.mockito.ArgumentMatcher;
  45
  46 import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos;
  47
  48 /**
  49  * Test case for {@link TableSpaceQuotaSnapshotNotifier}.
  50  */
  51 @Category(SmallTests.class)
  52 public class TestTableSpaceQuotaViolationNotifier {
  53
  54   @ClassRule
  55   public static final HBaseClassTestRule CLASS_RULE =
  56       HBaseClassTestRule.forClass(TestTableSpaceQuotaViolationNotifier.class);
  57
  58   private TableSpaceQuotaSnapshotNotifier notifier;
  59   private Connection conn;
  60
  61   @Before
  62   public void setup() throws Exception {
  63     notifier = new TableSpaceQuotaSnapshotNotifier();
  64     conn = mock(Connection.class);
  65     notifier.initialize(conn);
  66   }
  67
  68   @Test
  69   public void testToViolation() throws Exception {
  70     final TableName tn = TableName.valueOf("inviolation");
  71     final SpaceQuotaSnapshot snapshot = new SpaceQuotaSnapshot(
  72         new SpaceQuotaStatus(SpaceViolationPolicy.NO_INSERTS), 1024L, 512L);
  73     final Table quotaTable = mock(Table.class);
  74     when(conn.getTable(QuotaTableUtil.QUOTA_TABLE_NAME)).thenReturn(quotaTable);
  75
  76     final Put expectedPut = new Put(Bytes.toBytes("t." + tn.getNameAsString()));
  77     final QuotaProtos.SpaceQuotaSnapshot protoQuota = QuotaProtos.SpaceQuotaSnapshot.newBuilder()
  78         .setQuotaStatus(QuotaProtos.SpaceQuotaStatus.newBuilder().setInViolation(true)
  79         .setViolationPolicy(QuotaProtos.SpaceViolationPolicy.NO_INSERTS))
  80         .setQuotaLimit(512L)
  81         .setQuotaUsage(1024L)
  82         .build();
  83     expectedPut.addColumn(Bytes.toBytes("u"), Bytes.toBytes("p"), protoQuota.toByteArray());
  84
  85     notifier.transitionTable(tn, snapshot);
  86
  87     verify(quotaTable).put(argThat(new SingleCellMutationMatcher<Put>(expectedPut)));
  88   }
  89
  90   /**
  91    * Quick hack to verify a Mutation with one column.
  92    */
  93   final private static class SingleCellMutationMatcher<T> implements ArgumentMatcher<T> {
  94     private final Mutation expected;
  95
  96     private SingleCellMutationMatcher(Mutation expected) {
  97       this.expected = expected;
  98     }
  99
 100     @Override
 101     public boolean matches(T argument) {
 102       if (!expected.getClass().isAssignableFrom(argument.getClass())) {
 103         return false;
 104       }
 105       Mutation actual = (Mutation) argument;
 106       if (!Arrays.equals(expected.getRow(), actual.getRow())) {
 107         return false;
 108       }
 109       if (expected.size() != actual.size()) {
 110         return false;
 111       }
 112       NavigableMap<byte[],List<Cell>> expectedCells = expected.getFamilyCellMap();
 113       NavigableMap<byte[],List<Cell>> actualCells = actual.getFamilyCellMap();
 114       Entry<byte[],List<Cell>> expectedEntry = expectedCells.entrySet().iterator().next();
 115       Entry<byte[],List<Cell>> actualEntry = actualCells.entrySet().iterator().next();
 116       if (!Arrays.equals(expectedEntry.getKey(), actualEntry.getKey())) {
 117         return false;
 118       }
 119       return Objects.equals(expectedEntry.getValue(), actualEntry.getValue());
 120     }
 121   }
 122 }