bookkeeper-server/src/test/java/org/apache/bookkeeper/meta/MetadataDriversTest.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 org.apache.bookkeeper.meta;

import static org.apache.bookkeeper.meta.MetadataDrivers.BK_METADATA_BOOKIE_DRIVERS_PROPERTY;
import static org.apache.bookkeeper.meta.MetadataDrivers.BK_METADATA_CLIENT_DRIVERS_PROPERTY;
import static org.apache.bookkeeper.meta.MetadataDrivers.ZK_BOOKIE_DRIVER_CLASS;
import static org.apache.bookkeeper.meta.MetadataDrivers.ZK_CLIENT_DRIVER_CLASS;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import com.google.common.collect.Maps;
import java.net.URI;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ScheduledExecutorService;
import org.apache.bookkeeper.conf.ClientConfiguration;
import org.apache.bookkeeper.conf.ServerConfiguration;
import org.apache.bookkeeper.discover.RegistrationClient;
import org.apache.bookkeeper.discover.RegistrationManager;
import org.apache.bookkeeper.meta.MetadataDrivers.MetadataBookieDriverInfo;
import org.apache.bookkeeper.meta.MetadataDrivers.MetadataClientDriverInfo;
import org.apache.bookkeeper.meta.exceptions.MetadataException;
import org.apache.bookkeeper.stats.StatsLogger;
import org.apache.commons.lang3.StringUtils;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

/**
 * Unit test of {@link MetadataDrivers}.
 */
public class MetadataDriversTest {

    abstract static class TestClientDriver implements MetadataClientDriver {

        @Override
        public MetadataClientDriver initialize(ClientConfiguration conf,
                                               ScheduledExecutorService scheduler,
                                               StatsLogger statsLogger,
                                               Optional<Object> ctx) throws MetadataException {
            return this;
        }

        @Override
        public RegistrationClient getRegistrationClient() {
            return mock(RegistrationClient.class);
        }

        @Override
        public LedgerManagerFactory getLedgerManagerFactory() throws MetadataException {
            return mock(LedgerManagerFactory.class);
        }

        @Override
        public LayoutManager getLayoutManager() {
            return mock(LayoutManager.class);
        }

        @Override
        public void close() {
        }

        @Override
        public void setSessionStateListener(SessionStateListener sessionStateListener) {
        }
    }

    static class ClientDriver1 extends TestClientDriver {

        @Override
        public String getScheme() {
            return "driver1";
        }


    }

    static class ClientDriver2 extends TestClientDriver {

        @Override
        public String getScheme() {
            return "driver2";
        }

    }

    abstract static class TestBookieDriver implements MetadataBookieDriver {
        @Override
        public MetadataBookieDriver initialize(ServerConfiguration conf,
                                               StatsLogger statsLogger) throws MetadataException {
            return this;
        }

        @Override
        public RegistrationManager createRegistrationManager() {
            return mock(RegistrationManager.class);
        }

        @Override
        public LedgerManagerFactory getLedgerManagerFactory() throws MetadataException {
            return mock(LedgerManagerFactory.class);
        }

        @Override
        public LayoutManager getLayoutManager() {
            return mock(LayoutManager.class);
        }

        @Override
        public void close() {

        }
    }

    static class BookieDriver1 extends TestBookieDriver {

        @Override
        public String getScheme() {
            return "driver1";
        }

    }

    static class BookieDriver2 extends TestBookieDriver {

        @Override
        public String getScheme() {
            return "driver2";
        }

    }

    private Map<String, MetadataClientDriverInfo> savedClientDrivers;
    private Map<String, MetadataBookieDriverInfo> savedBookieDrivers;

    @Before
    public void setup() {
        savedClientDrivers = Maps.newHashMap(MetadataDrivers.getClientDrivers());
        savedBookieDrivers = Maps.newHashMap(MetadataDrivers.getBookieDrivers());
    }

    @After
    public void teardown() {
        MetadataDrivers.getClientDrivers().clear();
        MetadataDrivers.getClientDrivers().putAll(savedClientDrivers);
        MetadataDrivers.getBookieDrivers().clear();
        MetadataDrivers.getBookieDrivers().putAll(savedBookieDrivers);
    }

    @Test
    public void testDefaultDrivers() {
        MetadataClientDriver clientDriver = MetadataDrivers.getClientDriver("zk");
        assertEquals(
            ZK_CLIENT_DRIVER_CLASS,
            clientDriver.getClass().getName());
        clientDriver = MetadataDrivers.getClientDriver(URI.create("zk+hierarchical://127.0.0.1/ledgers"));
        assertEquals(
            ZK_CLIENT_DRIVER_CLASS,
            clientDriver.getClass().getName());

        MetadataBookieDriver bookieDriver = MetadataDrivers.getBookieDriver("zk");
        assertEquals(
            ZK_BOOKIE_DRIVER_CLASS,
            bookieDriver.getClass().getName());
        bookieDriver = MetadataDrivers.getBookieDriver(URI.create("zk+hierarchical://127.0.0.1/ledgers"));
        assertEquals(
            ZK_BOOKIE_DRIVER_CLASS,
            bookieDriver.getClass().getName());
    }

    @Test(expected = NullPointerException.class)
    public void testClientDriverNullScheme() {
        MetadataDrivers.getClientDriver((String) null);
    }

    @Test(expected = NullPointerException.class)
    public void testBookieDriverNullScheme() {
        MetadataDrivers.getBookieDriver((String) null);
    }

    @Test(expected = NullPointerException.class)
    public void testClientDriverNullURI() {
        MetadataDrivers.getClientDriver((URI) null);
    }

    @Test(expected = NullPointerException.class)
    public void testBookieDriverNullURI() {
        MetadataDrivers.getBookieDriver((URI) null);
    }

    @Test(expected = IllegalArgumentException.class)
    public void testClientDriverUnknownScheme() {
        MetadataDrivers.getClientDriver("unknown");
    }

    @Test(expected = IllegalArgumentException.class)
    public void testBookieDriverUnknownScheme() {
        MetadataDrivers.getBookieDriver("unknown");
    }

    @Test(expected = IllegalArgumentException.class)
    public void testClientDriverUnknownSchemeURI() {
        MetadataDrivers.getClientDriver(URI.create("unknown://"));
    }

    @Test(expected = IllegalArgumentException.class)
    public void testBookieDriverUnknownSchemeURI() {
        MetadataDrivers.getBookieDriver(URI.create("unknown://"));
    }

    @Test(expected = NullPointerException.class)
    public void testClientDriverNullSchemeURI() {
        MetadataDrivers.getClientDriver(URI.create("//127.0.0.1/ledgers"));
    }

    @Test(expected = NullPointerException.class)
    public void testBookieDriverNullSchemeURI() {
        MetadataDrivers.getBookieDriver(URI.create("//127.0.0.1/ledgers"));
    }

    @Test
    public void testClientDriverLowerUpperCasedSchemes() {
        String[] schemes = new String[] {
            "zk", "Zk", "zK", "ZK"
        };
        for (String scheme : schemes) {
            MetadataClientDriver clientDriver = MetadataDrivers.getClientDriver(scheme);
            assertEquals(
                ZK_CLIENT_DRIVER_CLASS,
                clientDriver.getClass().getName());
        }
    }

    @Test
    public void testBookieDriverLowerUpperCasedSchemes() {
        String[] schemes = new String[] {
            "zk", "Zk", "zK", "ZK"
        };
        for (String scheme : schemes) {
            MetadataBookieDriver bookieDriver = MetadataDrivers.getBookieDriver(scheme);
            assertEquals(
                ZK_BOOKIE_DRIVER_CLASS,
                bookieDriver.getClass().getName());
        }
    }

    @Test
    public void testRegisterClientDriver() throws Exception {
        MetadataClientDriver clientDriver = mock(MetadataClientDriver.class);
        when(clientDriver.getScheme()).thenReturn("testdriver");

        try {
            MetadataDrivers.getClientDriver(clientDriver.getScheme());
            fail("Should fail to get client driver if it is not registered");
        } catch (IllegalArgumentException iae) {
            // expected
        }

        MetadataDrivers.registerClientDriver(clientDriver.getScheme(), clientDriver.getClass());
        MetadataClientDriver driver = MetadataDrivers.getClientDriver(clientDriver.getScheme());
        assertEquals(clientDriver.getClass(), driver.getClass());
    }

    @Test
    public void testRegisterBookieDriver() throws Exception {
        MetadataBookieDriver bookieDriver = mock(MetadataBookieDriver.class);
        when(bookieDriver.getScheme()).thenReturn("testdriver");

        try {
            MetadataDrivers.getBookieDriver(bookieDriver.getScheme());
            fail("Should fail to get bookie driver if it is not registered");
        } catch (IllegalArgumentException iae) {
            // expected
        }

        MetadataDrivers.registerBookieDriver(bookieDriver.getScheme(), bookieDriver.getClass());
        MetadataBookieDriver driver = MetadataDrivers.getBookieDriver(bookieDriver.getScheme());
        assertEquals(bookieDriver.getClass(), driver.getClass());
    }

    @Test
    public void testLoadClientDriverFromSystemProperty() throws Exception {
        String saveDriversStr = System.getProperty(BK_METADATA_CLIENT_DRIVERS_PROPERTY);
        try {
            System.setProperty(
                BK_METADATA_CLIENT_DRIVERS_PROPERTY,
                StringUtils.join(new String[] {
                    ClientDriver1.class.getName(),
                    ClientDriver2.class.getName()
                }, ':'));

            MetadataDrivers.loadInitialDrivers();

            MetadataClientDriver loadedDriver1 = MetadataDrivers.getClientDriver("driver1");
            assertEquals(ClientDriver1.class, loadedDriver1.getClass());
            MetadataClientDriver loadedDriver2 = MetadataDrivers.getClientDriver("driver2");
            assertEquals(ClientDriver2.class, loadedDriver2.getClass());
        } finally {
            if (null != saveDriversStr) {
                System.setProperty(BK_METADATA_CLIENT_DRIVERS_PROPERTY, saveDriversStr);
            } else {
                System.clearProperty(BK_METADATA_CLIENT_DRIVERS_PROPERTY);
            }
        }
    }

    @Test
    public void testLoadBookieDriverFromSystemProperty() throws Exception {
        String saveDriversStr = System.getProperty(BK_METADATA_BOOKIE_DRIVERS_PROPERTY);
        try {
            System.setProperty(
                BK_METADATA_BOOKIE_DRIVERS_PROPERTY,
                StringUtils.join(new String[] {
                    BookieDriver1.class.getName(),
                    BookieDriver2.class.getName()
                }, ':'));

            MetadataDrivers.loadInitialDrivers();

            MetadataBookieDriver loadedDriver1 = MetadataDrivers.getBookieDriver("driver1");
            assertEquals(BookieDriver1.class, loadedDriver1.getClass());
            MetadataBookieDriver loadedDriver2 = MetadataDrivers.getBookieDriver("driver2");
            assertEquals(BookieDriver2.class, loadedDriver2.getClass());
        } finally {
            if (null != saveDriversStr) {
                System.setProperty(BK_METADATA_BOOKIE_DRIVERS_PROPERTY, saveDriversStr);
            } else {
                System.clearProperty(BK_METADATA_BOOKIE_DRIVERS_PROPERTY);
            }
        }
    }

}