HiveNotificationFetcher implements AutoCloseable { public List<NotificationEvent> fetchNotifications(long lastEventId) throws Exception { return fetchNotifications(lastEventId, Integer.MAX_VALUE); } HiveNotificationFetcher(SentryStoreInterface sentryStore, HiveConnectionFactory hmsConnectionFactory); List<NotificationEvent> fetchNotifications(long lastEventId); List<NotificationEvent> fetchNotifications(long lastEventId, int maxEvents); long getCurrentNotificationId(); @Override void close();  }
@Test public void testGetEmptyNotificationsWhenHmsReturnsANullResponse() throws Exception { SentryStore store = Mockito.mock(SentryStore.class); HiveConnectionFactory hmsConnection = Mockito.mock(HiveConnectionFactory.class); HiveMetaStoreClient hmsClient = Mockito.mock(HiveMetaStoreClient.class); Mockito.when(hmsConnection.connect()).thenReturn(new HMSClient(hmsClient)); try (HiveNotificationFetcher fetcher = new HiveNotificationFetcher(store, hmsConnection)) { List<NotificationEvent> events; Mockito.when(hmsClient.getNextNotification(0, Integer.MAX_VALUE, null)) .thenReturn(null); events = fetcher.fetchNotifications(0); assertTrue(events.isEmpty()); } }
@Test public void testGetEmptyNotificationsWhenHmsReturnsEmptyEvents() throws Exception { SentryStore store = Mockito.mock(SentryStore.class); HiveConnectionFactory hmsConnection = Mockito.mock(HiveConnectionFactory.class); HiveMetaStoreClient hmsClient = Mockito.mock(HiveMetaStoreClient.class); Mockito.when(hmsConnection.connect()).thenReturn(new HMSClient(hmsClient)); try (HiveNotificationFetcher fetcher = new HiveNotificationFetcher(store, hmsConnection)) { List<NotificationEvent> events; Mockito.when(hmsClient.getNextNotification(0, Integer.MAX_VALUE, null)) .thenReturn(new NotificationEventResponse(Collections.<NotificationEvent>emptyList())); events = fetcher.fetchNotifications(0); assertTrue(events.isEmpty()); } }
@Test public void testGetAllNotificationsReturnedByHms() throws Exception { SentryStore store = Mockito.mock(SentryStore.class); HiveConnectionFactory hmsConnection = Mockito.mock(HiveConnectionFactory.class); HiveMetaStoreClient hmsClient = Mockito.mock(HiveMetaStoreClient.class); Mockito.when(hmsConnection.connect()).thenReturn(new HMSClient(hmsClient)); try (HiveNotificationFetcher fetcher = new HiveNotificationFetcher(store, hmsConnection)) { List<NotificationEvent> events; Mockito.when(hmsClient.getNextNotification(0, Integer.MAX_VALUE, null)) .thenReturn(new NotificationEventResponse( Arrays.<NotificationEvent>asList( new NotificationEvent(1L, 0, "CREATE_DATABASE", ""), new NotificationEvent(2L, 0, "CREATE_TABLE", "") ) )); events = fetcher.fetchNotifications(0); assertEquals(2, events.size()); assertEquals(1, events.get(0).getEventId()); assertEquals("CREATE_DATABASE", events.get(0).getEventType()); assertEquals(2, events.get(1).getEventId()); assertEquals("CREATE_TABLE", events.get(1).getEventType()); } }
@Test public void testGetDuplicatedEventsAndFilterEventsAlreadySeen() throws Exception { final SentryStore store = Mockito.mock(SentryStore.class); HiveConnectionFactory hmsConnection = Mockito.mock(HiveConnectionFactory.class); HiveMetaStoreClient hmsClient = Mockito.mock(HiveMetaStoreClient.class); Mockito.when(hmsConnection.connect()).thenReturn(new HMSClient(hmsClient)); try (HiveNotificationFetcher fetcher = new HiveNotificationFetcher(store, hmsConnection)) { List<NotificationEvent> events; Mockito.when(hmsClient.getNextNotification(Mockito.eq(0L), Mockito.eq(Integer.MAX_VALUE), (NotificationFilter) Mockito.notNull())).thenAnswer(new Answer<NotificationEventResponse>() { @Override public NotificationEventResponse answer(InvocationOnMock invocation) throws Throwable { NotificationFilter filter = (NotificationFilter) invocation.getArguments()[2]; NotificationEventResponse response = new NotificationEventResponse(); List<NotificationEvent> events = Arrays.<NotificationEvent>asList( new NotificationEvent(1L, 0, "CREATE_DATABASE", ""), new NotificationEvent(1L, 0, "CREATE_TABLE", ""), new NotificationEvent(2L, 0, "ALTER_TABLE", "") ); for (NotificationEvent event : events) { String hash = UniquePathsUpdate.sha1(event); if (event.getEventType().equals("CREATE_DATABASE")) { Mockito.when(store.isNotificationProcessed(Mockito.eq(hash))).thenReturn(true); } else { Mockito.when(store.isNotificationProcessed(Mockito.eq(hash))).thenReturn(false); } if (filter.accept(event)) { response.addToEvents(event); } } return response; } }); events = fetcher.fetchNotifications(1); assertEquals(2, events.size()); assertEquals(1, events.get(0).getEventId()); assertEquals("CREATE_TABLE", events.get(0).getEventType()); assertEquals(2, events.get(1).getEventId()); assertEquals("ALTER_TABLE", events.get(1).getEventType()); } }
@Test public void testPartialFetchesFromHMS() throws Exception { SentryStore store = Mockito.mock(SentryStore.class); HiveConnectionFactory hmsConnection = Mockito.mock(HiveConnectionFactory.class); HiveMetaStoreClient hmsClient = Mockito.mock(HiveMetaStoreClient.class); Mockito.when(hmsConnection.connect()).thenReturn(new HMSClient(hmsClient)); try (HiveNotificationFetcher fetcher = new HiveNotificationFetcher(store, hmsConnection)) { List<NotificationEvent> events; Mockito.when(hmsClient.getNextNotification(0, Integer.MAX_VALUE, null)) .thenReturn(new NotificationEventResponse( Arrays.<NotificationEvent>asList( new NotificationEvent(1L, 0, "CREATE_DATABASE", ""), new NotificationEvent(2L, 0, "CREATE_TABLE", "") ) )); Mockito.when(hmsClient.getNextNotification(0, 1, null)) .thenReturn(new NotificationEventResponse( Arrays.<NotificationEvent>asList( new NotificationEvent(1L, 0, "CREATE_DATABASE", "") ) )); int sentryHMSFetchSize = 1; events = fetcher.fetchNotifications(0, sentryHMSFetchSize); assertEquals(1, events.size()); assertEquals(1, events.get(0).getEventId()); assertEquals("CREATE_DATABASE", events.get(0).getEventType()); } }