content: /*
   * Copyright 2010 Proofpoint, Inc.
   *
   * Licensed 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 io.airlift.discovery.server;
  
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.ImmutableSet;
  import com.google.inject.Binder;
  import com.google.inject.Injector;
  import com.google.inject.Key;
  import com.google.inject.Module;
  import io.airlift.bootstrap.Bootstrap;
  import io.airlift.discovery.client.DiscoveryAnnouncementClient;
  import io.airlift.discovery.client.DiscoveryModule;
  import io.airlift.discovery.client.ServiceAnnouncement;
  import io.airlift.discovery.client.ServiceDescriptor;
  import io.airlift.discovery.client.ServiceSelector;
  import io.airlift.http.client.HttpClient;
  import io.airlift.http.client.Request;
  import io.airlift.http.client.jetty.JettyHttpClient;
  import io.airlift.http.server.testing.TestingHttpServer;
  import io.airlift.http.server.testing.TestingHttpServerModule;
  import io.airlift.jaxrs.JaxrsModule;
  import io.airlift.json.JsonModule;
  import io.airlift.node.NodeInfo;
  import io.airlift.node.testing.TestingNodeModule;
  import jakarta.ws.rs.core.HttpHeaders;
  import jakarta.ws.rs.core.MediaType;
  import org.iq80.leveldb.util.FileUtils;
  import org.testng.annotations.AfterMethod;
  import org.testng.annotations.BeforeMethod;
  import org.testng.annotations.Test;
  import org.weakref.jmx.guice.MBeanModule;
  import org.weakref.jmx.testing.TestingMBeanServer;
  
  import javax.management.MBeanServer;
  
  import java.io.File;
  import java.net.URI;
  import java.util.List;
  import java.util.Map;
  
  import static io.airlift.discovery.client.DiscoveryBinder.discoveryBinder;
  import static io.airlift.discovery.client.ServiceTypes.serviceType;
  import static io.airlift.http.client.FullJsonResponseHandler.JsonResponse;
  import static io.airlift.http.client.FullJsonResponseHandler.createFullJsonResponseHandler;
  import static io.airlift.http.client.JsonBodyGenerator.jsonBodyGenerator;
  import static io.airlift.http.client.Request.Builder.prepareDelete;
  import static io.airlift.http.client.Request.Builder.preparePost;
  import static io.airlift.http.client.StatusResponseHandler.StatusResponse;
  import static io.airlift.http.client.StatusResponseHandler.createStatusResponseHandler;
  import static io.airlift.json.JsonCodec.jsonCodec;
  import static io.airlift.json.JsonCodec.mapJsonCodec;
  import static jakarta.ws.rs.core.Response.Status;
  import static java.nio.file.Files.createTempDirectory;
  import static org.testng.Assert.assertEquals;
  import static org.testng.Assert.assertNotNull;
  import static org.testng.Assert.assertNull;
  import static org.testng.Assert.assertTrue;
  
  public class TestDiscoveryServer
  {
      private TestingHttpServer server;
      private File tempDir;
  
      @BeforeMethod
      public void setup()
              throws Exception
      {
          tempDir = createTempDirectory(null).toFile();
  
          // start server
          Map<String, String> serverProperties = ImmutableMap.<String, String>builder()
                  .put("static.db.location", tempDir.getAbsolutePath())
                  .put("discovery.store-cache-ttl", "0s")
                  .build();
  
          Bootstrap bootstrap = new Bootstrap(
                  new MBeanModule(),
                  new TestingNodeModule("testing"),
                  new TestingHttpServerModule(),
                  new JsonModule(),
                  new JaxrsModule(),
                  new DiscoveryServerModule(),
                  new DiscoveryModule(),
                  new Module()
                  {
                      @Override
                      public void configure(Binder binder)
                      {
                          // TODO: use testing mbean server
                          binder.bind(MBeanServer.class).toInstance(new TestingMBeanServer());
                      }
                  });
  
          Injector serverInjector = bootstrap
                  .doNotInitializeLogging()
                  .setRequiredConfigurationProperties(serverProperties)
                  .initialize();
  
          server = serverInjector.getInstance(TestingHttpServer.class);
          server.start();
      }
  
      @AfterMethod
      public void teardown()
              throws Exception
      {
          server.stop();
          FileUtils.deleteRecursively(tempDir);
      }
  
      @Test
      public void testDynamicAnnouncement()
              throws Exception
      {
          // publish announcement
          Map<String, String> announcerProperties = ImmutableMap.<String, String>builder()
                  .put("discovery.uri", server.getBaseUrl().toString())
                  .build();
  
          Bootstrap bootstrap = new Bootstrap(
                  new TestingNodeModule("testing", "red"),
                  new JsonModule(),
                  new DiscoveryModule());
  
          Injector announcerInjector = bootstrap
                  .doNotInitializeLogging()
                  .setRequiredConfigurationProperties(announcerProperties)
                  .initialize();
  
          ServiceAnnouncement announcement = ServiceAnnouncement.serviceAnnouncement("apple")
                  .addProperties(ImmutableMap.of("key", "value"))
                  .build();
  
          DiscoveryAnnouncementClient client = announcerInjector.getInstance(DiscoveryAnnouncementClient.class);
          client.announce(ImmutableSet.of(announcement)).get();
  
          NodeInfo announcerNodeInfo = announcerInjector.getInstance(NodeInfo.class);
  
          List<ServiceDescriptor> services = selectorFor("apple", "red").selectAllServices();
          assertEquals(services.size(), 1);
  
          ServiceDescriptor service = services.get(0);
          assertNotNull(service.getId());
          assertEquals(service.getNodeId(), announcerNodeInfo.getNodeId());
          assertEquals(service.getLocation(), announcerNodeInfo.getLocation());
          assertEquals(service.getPool(), announcerNodeInfo.getPool());
          assertEquals(service.getProperties(), announcement.getProperties());
  
          // ensure that service is no longer visible
          client.unannounce().get();
  
          assertTrue(selectorFor("apple", "red").selectAllServices().isEmpty());
      }
  
      @Test
      public void testStaticAnnouncement()
              throws Exception
      {
          // create static announcement
          Map<String, Object> announcement = ImmutableMap.<String, Object>builder()
                  .put("environment", "testing")
                  .put("type", "apple")
                  .put("pool", "red")
                  .put("location", "/a/b/c")
                  .put("properties", ImmutableMap.of("http", "http://host"))
                  .build();
  
          HttpClient client = new JettyHttpClient();
  
          Request request = preparePost()
                  .setUri(uriFor("/v1/announcement/static"))
                  .setHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON)
                  .setBodyGenerator(jsonBodyGenerator(jsonCodec(Object.class), announcement))
                  .build();
          JsonResponse<Map<String, Object>> createResponse = client.execute(request, createFullJsonResponseHandler(mapJsonCodec(String.class, Object.class)));
  
          assertEquals(createResponse.getStatusCode(), Status.CREATED.getStatusCode());
          String id = createResponse.getValue().get("id").toString();
  
          List<ServiceDescriptor> services = selectorFor("apple", "red").selectAllServices();
          assertEquals(services.size(), 1);
  
          ServiceDescriptor service = services.get(0);
          assertEquals(service.getId().toString(), id);
          assertNull(service.getNodeId());
          assertEquals(service.getLocation(), announcement.get("location"));
          assertEquals(service.getPool(), announcement.get("pool"));
          assertEquals(service.getProperties(), announcement.get("properties"));
  
          // remove announcement
          request = prepareDelete().setUri(uriFor("/v1/announcement/static/" + id)).build();
          StatusResponse deleteResponse = client.execute(request, createStatusResponseHandler());
  
          assertEquals(deleteResponse.getStatusCode(), Status.NO_CONTENT.getStatusCode());
  
          // ensure announcement is gone
          assertTrue(selectorFor("apple", "red").selectAllServices().isEmpty());
      }
  
      private ServiceSelector selectorFor(String type, String pool)
              throws Exception
      {
          Map<String, String> clientProperties = ImmutableMap.<String, String>builder()
                  .put("discovery.uri", server.getBaseUrl().toString())
                  .put("discovery.apple.pool", "red")
                  .build();
  
          Bootstrap bootstrap = new Bootstrap(
                  new TestingNodeModule("testing"),
                  new JsonModule(),
                  new DiscoveryModule(),
                  new Module()
                  {
                      @Override
                      public void configure(Binder binder)
                      {
                          discoveryBinder(binder).bindSelector("apple");
                      }
                  });
  
          Injector clientInjector = bootstrap
                  .doNotInitializeLogging()
                  .setRequiredConfigurationProperties(clientProperties)
                  .initialize();
  
          return clientInjector.getInstance(Key.get(ServiceSelector.class, serviceType("apple")));
      }
  
      private URI uriFor(String path)
      {
          return server.getBaseUrl().resolve(path);
      }
  }
  