renamed and aded test

ExpediaGroup · Apr 12, 2024 · 070ce28 · 070ce28
1 parent 91fa874
commit 070ce28
Show file tree

Hide file tree

Showing 10 changed files with 51 additions and 32 deletions.
diff --git a/...ava/com/hotels/bdp/waggledance/client/CloseableThriftHiveMetastoreIfaceClientFactory.java b/...ava/com/hotels/bdp/waggledance/client/CloseableThriftHiveMetastoreIfaceClientFactory.java
@@ -25,11 +25,8 @@
 
 import org.apache.hadoop.hive.conf.HiveConf.ConfVars;
 import org.apache.hadoop.hive.metastore.api.MetaException;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import com.amazonaws.glue.catalog.metastore.AWSCatalogMetastoreClient;
-
 import com.hotels.bdp.waggledance.api.WaggleDanceException;
 import com.hotels.bdp.waggledance.api.model.AbstractMetaStore;
 import com.hotels.bdp.waggledance.client.adapter.MetastoreIfaceAdapter;

diff --git a/...e-core/src/main/java/com/hotels/bdp/waggledance/client/DefaultMetaStoreClientFactory.java b/...e-core/src/main/java/com/hotels/bdp/waggledance/client/DefaultMetaStoreClientFactory.java
@@ -162,7 +162,6 @@ private SaslMetastoreClientHander(
       this.clientManager = clientManager;
     }
 
-    @SuppressWarnings("unchecked")
     @Override
     public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
       try {

diff --git a/...-dance-extensions/src/main/java/com/hotels/bdp/waggledance/extensions/ExtensionBeans.java b/...-dance-extensions/src/main/java/com/hotels/bdp/waggledance/extensions/ExtensionBeans.java
@@ -29,24 +29,24 @@
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
-import io.github.bucket4j.distributed.ExpirationAfterWriteStrategy;
-import io.github.bucket4j.distributed.serialization.Mapper;
-import io.github.bucket4j.redis.redisson.cas.RedissonBasedProxyManager;
-
 import com.hotels.bdp.waggledance.client.ThriftClientFactory;
-import com.hotels.bdp.waggledance.extensions.client.ratelimit.BucketBandwithProvider;
+import com.hotels.bdp.waggledance.extensions.client.ratelimit.BucketBandwidthProvider;
 import com.hotels.bdp.waggledance.extensions.client.ratelimit.BucketService;
 import com.hotels.bdp.waggledance.extensions.client.ratelimit.RateLimitingClientFactory;
 import com.hotels.bdp.waggledance.extensions.client.ratelimit.RefillType;
 import com.hotels.bdp.waggledance.extensions.client.ratelimit.memory.InMemoryBucketService;
 import com.hotels.bdp.waggledance.extensions.client.ratelimit.redis.RedisBucketService;
 
+import io.github.bucket4j.distributed.ExpirationAfterWriteStrategy;
+import io.github.bucket4j.distributed.serialization.Mapper;
+import io.github.bucket4j.redis.redisson.cas.RedissonBasedProxyManager;
+
 @Configuration
 @ConditionalOnProperty(name = "waggledance.extensions.ratelimit.enabled", havingValue = "true")
 public class ExtensionBeans {
 
   private static final String STORAGE_MEMORY = "MEMORY";
-  private static final String STORAGE_REDIS= "REDIS";
+  private static final String STORAGE_REDIS = "REDIS";
 
   @Bean
   public ThriftClientFactory thriftClientFactory(
@@ -57,16 +57,16 @@ public ThriftClientFactory thriftClientFactory(
 
   @ConditionalOnProperty(name = "waggledance.extensions.ratelimit.storage", havingValue = STORAGE_MEMORY)
   @Bean
-  public BucketService inMemoryBucketService(BucketBandwithProvider bucketBandwithProvider) {
-    return new InMemoryBucketService(bucketBandwithProvider);
+  public BucketService inMemoryBucketService(BucketBandwidthProvider bucketBandwidthProvider) {
+    return new InMemoryBucketService(bucketBandwidthProvider);
   }
 
   @ConditionalOnProperty(name = "waggledance.extensions.ratelimit.storage", havingValue = STORAGE_REDIS)
   @Bean
   public BucketService redisBucketService(
-      BucketBandwithProvider bucketBandwithProvider,
+      BucketBandwidthProvider bucketBandwidthProvider,
       RedissonBasedProxyManager<String> redissonBasedProxyManager) {
-    return new RedisBucketService(bucketBandwithProvider, redissonBasedProxyManager);
+    return new RedisBucketService(bucketBandwidthProvider, redissonBasedProxyManager);
   }
 
   @ConditionalOnProperty(name = "waggledance.extensions.ratelimit.storage", havingValue = STORAGE_REDIS)
@@ -88,7 +88,7 @@ public RedissonBasedProxyManager<String> redissonBasedProxyManager(
   }
 
   @Bean
-  public BucketBandwithProvider bucketBandwithProvider(
+  public BucketBandwidthProvider bucketBandwidthProvider(
       @Value("${waggledance.extensions.ratelimit.capacity:2000}") long capacity,
       @Value("${waggledance.extensions.ratelimit.tokensPerMinute:1000}") long tokensPerMinute,
       @Value("${waggledance.extensions.ratelimit.refillType:GREEDY}") RefillType refillType) {

diff --git a/...ent/ratelimit/BucketBandwithProvider.java → ...nt/ratelimit/BucketBandwidthProvider.java b/...ent/ratelimit/BucketBandwithProvider.java → ...nt/ratelimit/BucketBandwidthProvider.java
diff --git a/.../java/com/hotels/bdp/waggledance/extensions/client/ratelimit/GreedyBandwidthProvider.java b/.../java/com/hotels/bdp/waggledance/extensions/client/ratelimit/GreedyBandwidthProvider.java
@@ -19,7 +19,7 @@
 
 import io.github.bucket4j.Bandwidth;
 
-public class GreedyBandwidthProvider implements BucketBandwithProvider {
+public class GreedyBandwidthProvider implements BucketBandwidthProvider {
 
   private final long capacity;
   private final long tokensPerMinute;
@@ -30,7 +30,7 @@ public GreedyBandwidthProvider(long capacity, long tokensPerMinute) {
   }
 
   @Override
-  public Bandwidth getBandwith() {
+  public Bandwidth getBandwidth() {
     Bandwidth limit = Bandwidth
         .builder()
         .capacity(capacity)

diff --git a/...a/com/hotels/bdp/waggledance/extensions/client/ratelimit/IntervallyBandwidthProvider.java b/...a/com/hotels/bdp/waggledance/extensions/client/ratelimit/IntervallyBandwidthProvider.java
@@ -19,7 +19,7 @@
 
 import io.github.bucket4j.Bandwidth;
 
-public class IntervallyBandwidthProvider implements BucketBandwithProvider {
+public class IntervallyBandwidthProvider implements BucketBandwidthProvider {
 
   private long capacity;
   private long tokensPerMinute;
@@ -30,7 +30,7 @@ public IntervallyBandwidthProvider(long capacity, long tokensPerMinute) {
   }
 
   @Override
-  public Bandwidth getBandwith() {
+  public Bandwidth getBandwidth() {
     Bandwidth limit = Bandwidth
         .builder()
         .capacity(capacity)

diff --git a/...ions/src/main/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RefillType.java b/...ions/src/main/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RefillType.java
@@ -18,16 +18,16 @@
 public enum RefillType {
   GREEDY {
     @Override
-    public BucketBandwithProvider createBandwidthProvider(long capacity, long tokensPerMinute) {
+    public BucketBandwidthProvider createBandwidthProvider(long capacity, long tokensPerMinute) {
       return new GreedyBandwidthProvider(capacity, tokensPerMinute);
     }
   },
   INTERVALLY {
     @Override
-    public BucketBandwithProvider createBandwidthProvider(long capacity, long tokensPerMinute) {
+    public BucketBandwidthProvider createBandwidthProvider(long capacity, long tokensPerMinute) {
       return new IntervallyBandwidthProvider(capacity, tokensPerMinute);
     }
   };
 
-  public abstract BucketBandwithProvider createBandwidthProvider(long capacity, long tokensPerMinute);
+  public abstract BucketBandwidthProvider createBandwidthProvider(long capacity, long tokensPerMinute);
 }
diff --git a/.../com/hotels/bdp/waggledance/extensions/client/ratelimit/memory/InMemoryBucketService.java b/.../com/hotels/bdp/waggledance/extensions/client/ratelimit/memory/InMemoryBucketService.java
@@ -20,7 +20,7 @@
 
 import io.github.bucket4j.Bucket;
 
-import com.hotels.bdp.waggledance.extensions.client.ratelimit.BucketBandwithProvider;
+import com.hotels.bdp.waggledance.extensions.client.ratelimit.BucketBandwidthProvider;
 import com.hotels.bdp.waggledance.extensions.client.ratelimit.BucketService;
 
 /**
@@ -29,15 +29,15 @@
  */
 public class InMemoryBucketService implements BucketService {
 
-  private final BucketBandwithProvider bucketBandwithProvider;
+  private final BucketBandwidthProvider bucketBandwidthProvider;
   private Map<String, Bucket> bucketsPerUser = new HashMap<>();
 
-  public InMemoryBucketService(BucketBandwithProvider bucketBandwithProvider) {
-    this.bucketBandwithProvider = bucketBandwithProvider;
+  public InMemoryBucketService(BucketBandwidthProvider bucketBandwidthProvider) {
+    this.bucketBandwidthProvider = bucketBandwidthProvider;
   }
 
   private Bucket createNewBucket() {
-    return Bucket.builder().addLimit(bucketBandwithProvider.getBandwith()).build();
+    return Bucket.builder().addLimit(bucketBandwidthProvider.getBandwidth()).build();
   }
 
   public Bucket getBucket(String key) {

diff --git a/...java/com/hotels/bdp/waggledance/extensions/client/ratelimit/redis/RedisBucketService.java b/...java/com/hotels/bdp/waggledance/extensions/client/ratelimit/redis/RedisBucketService.java
@@ -19,7 +19,7 @@
 import io.github.bucket4j.BucketConfiguration;
 import io.github.bucket4j.redis.redisson.cas.RedissonBasedProxyManager;
 
-import com.hotels.bdp.waggledance.extensions.client.ratelimit.BucketBandwithProvider;
+import com.hotels.bdp.waggledance.extensions.client.ratelimit.BucketBandwidthProvider;
 import com.hotels.bdp.waggledance.extensions.client.ratelimit.BucketService;
 
 /**
@@ -32,10 +32,10 @@ public class RedisBucketService implements BucketService {
   private final BucketConfiguration configuration;
 
   public RedisBucketService(
-      BucketBandwithProvider bucketBandwithProvider,
+      BucketBandwidthProvider bucketBandwidthProvider,
       RedissonBasedProxyManager<String> proxyManager) {
     this.proxyManager = proxyManager;
-    configuration = BucketConfiguration.builder().addLimit(bucketBandwithProvider.getBandwith()).build();
+    configuration = BucketConfiguration.builder().addLimit(bucketBandwidthProvider.getBandwidth()).build();
   }
 
   @Override

diff --git a/...hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandlerTest.java b/...hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandlerTest.java
@@ -15,14 +15,15 @@
  */
 package com.hotels.bdp.waggledance.extensions.client.ratelimit;
 
+import static com.hotels.bdp.waggledance.extensions.client.ratelimit.RateLimitingInvocationHandler.UNKNOWN_USER;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import static com.hotels.bdp.waggledance.extensions.client.ratelimit.RateLimitingInvocationHandler.UNKNOWN_USER;
-
 import org.apache.hadoop.hive.metastore.api.NoSuchObjectException;
 import org.apache.hadoop.hive.metastore.api.Table;
 import org.junit.Test;
@@ -71,6 +72,9 @@ public void testLimitDifferentUsers() throws Exception {
     } catch (WaggleDanceServerException e) {
       assertThat(e.getMessage(), is("[STATUS=429] Too many requests."));
     }
+
+    verify(client, times(3)).get_table("db", "table");
+    verify(client).set_ugi(USER, null);
   }
 
   @Test
@@ -103,6 +107,25 @@ public void testInvocationHandlerThrowsCause() throws Exception {
       assertThat(e.getMessage(), is("No such table"));
     }
   }
+
+  @Test
+  public void testIgnoredMethods() throws Exception {
+    when(thriftClientFactory.newInstance(metastore)).thenReturn(client);
+    CloseableThriftHiveMetastoreIface handlerProxy = new RateLimitingClientFactory(thriftClientFactory, bucketService)
+        .newInstance(metastore);
+
+    assertTokens(2, 2);
+    handlerProxy.set_ugi(USER, null);
+    handlerProxy.isOpen();
+    handlerProxy.flushCache();
+    handlerProxy.close();
+    assertTokens(2, 2);
+
+    verify(client).set_ugi(USER, null);
+    verify(client).isOpen();
+    verify(client).flushCache();
+    verify(client).close();
+  }
 
   private void assertTokens(long expectedUserTokenCount, long expectedUnknownUserTokenCount) {
     assertThat(bucketService.getBucket(USER).getAvailableTokens(), is(expectedUserTokenCount));