content: /*
   * Copyright 1999-2023 Alibaba Group Holding Ltd.
   *
   * 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 com.alibaba.nacos.test.ability;
  
  import com.alibaba.nacos.Nacos;
  import com.alibaba.nacos.api.NacosFactory;
  import com.alibaba.nacos.api.PropertyKeyConst;
  import com.alibaba.nacos.api.ability.constant.AbilityKey;
  import com.alibaba.nacos.api.ability.constant.AbilityStatus;
  import com.alibaba.nacos.api.config.ConfigService;
  import com.alibaba.nacos.api.config.remote.request.ConfigQueryRequest;
  import com.alibaba.nacos.api.config.remote.response.ConfigQueryResponse;
  import com.alibaba.nacos.api.exception.NacosException;
  import com.alibaba.nacos.api.naming.remote.request.NotifySubscriberRequest;
  import com.alibaba.nacos.api.remote.request.Request;
  import com.alibaba.nacos.api.remote.request.RequestMeta;
  import com.alibaba.nacos.api.remote.response.Response;
  import com.alibaba.nacos.common.ability.AbstractAbilityControlManager;
  import com.alibaba.nacos.common.ability.discover.NacosAbilityManagerHolder;
  import com.alibaba.nacos.common.remote.ConnectionType;
  import com.alibaba.nacos.common.remote.client.Connection;
  import com.alibaba.nacos.common.remote.client.RpcClient;
  import com.alibaba.nacos.common.remote.client.RpcClientFactory;
  import com.alibaba.nacos.common.remote.client.ServerListFactory;
  import com.alibaba.nacos.common.remote.client.ServerRequestHandler;
  import com.alibaba.nacos.common.remote.client.grpc.GrpcClient;
  import com.alibaba.nacos.core.remote.ConnectionManager;
  import com.alibaba.nacos.core.remote.RequestFilters;
  import com.alibaba.nacos.core.remote.RequestHandler;
  import com.alibaba.nacos.core.remote.RequestHandlerRegistry;
  import com.alibaba.nacos.test.ability.component.TestServerAbilityControlManager;
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.extension.ExtendWith;
  import org.springframework.boot.test.context.SpringBootTest;
  import org.springframework.boot.test.web.server.LocalServerPort;
  import org.springframework.test.context.junit.jupiter.SpringExtension;
  
  import javax.annotation.Resource;
  import java.lang.reflect.Field;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Properties;
  import java.util.UUID;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  @ExtendWith(SpringExtension.class)
  @SpringBootTest(classes = Nacos.class, properties = {
          "server.servlet.context-path=/nacos"}, webEnvironment = SpringBootTest.WebEnvironment.DEFINED_PORT)
  @SuppressWarnings("all")
  class AbilityDiscovery {
      
      @LocalServerPort
      private int port;
      
      @Resource
      private RequestHandlerRegistry requestHandlerRegistry;
      
      @Resource
      private RequestFilters filters;
      
      @Resource
      private ConnectionManager connectionManager;
      
      private RpcClient client;
      
      private RpcClient clusterClient;
      
      private ConfigService configService;
      
      private AbstractAbilityControlManager oldInstance;
      
      /**
       * test server judge client abilities
       */
      private volatile boolean serverSuccess = false;
      
      private volatile boolean clientSuccess = false;
      
      private volatile boolean clusterSuccess = false;
      
      private Field abstractAbilityControlManager;
      
      private Field registryHandlerFields;
      
      private Field serverReuqestHandlersField;
      
      private Field currentConnField;
      
      private Field setupRequestHandlerField;
      
      @BeforeEach
      void setup() throws NoSuchFieldException, IllegalAccessException, NacosException {
          // load class
          oldInstance = NacosAbilityManagerHolder.getInstance();
          
          // replace
          abstractAbilityControlManager = NacosAbilityManagerHolder.class.getDeclaredField("abstractAbilityControlManager");
          abstractAbilityControlManager.setAccessible(true);
          abstractAbilityControlManager.set(NacosAbilityManagerHolder.class, new TestServerAbilityControlManager());
          
          // get registry field
          registryHandlerFields = RequestHandlerRegistry.class.getDeclaredField("registryHandlers");
          registryHandlerFields.setAccessible(true);
          
          // currentConn
          currentConnField = RpcClient.class.getDeclaredField("currentConnection");
          currentConnField.setAccessible(true);
          
          // init config service
          Properties properties = new Properties();
          properties.put(PropertyKeyConst.SERVER_ADDR, "127.0.0.1:" + port);
          configService = NacosFactory.createConfigService(properties);
          
          // server request handler
          serverReuqestHandlersField = RpcClient.class.getDeclaredField("serverRequestHandlers");
          serverReuqestHandlersField.setAccessible(true);
          
          // setupRequestHandler
          setupRequestHandlerField = GrpcClient.class.getDeclaredField("setupRequestHandler");
          setupRequestHandlerField.setAccessible(true);
          
          // init client
          client = RpcClientFactory.createClient(UUID.randomUUID().toString(), ConnectionType.GRPC, new HashMap<>());
          client.serverListFactory(new ServerListFactory() {
              @Override
              public String genNextServer() {
                  return "127.0.0.1:" + port;
              }
              
              @Override
              public String getCurrentServer() {
                  return "127.0.0.1:" + port;
              }
              
              @Override
              public List<String> getServerList() {
                  return Collections.singletonList("127.0.0.1:" + port);
              }
          });
          // connect to server
          client.start();
          
          clusterClient = RpcClientFactory.createClusterClient(UUID.randomUUID().toString(), ConnectionType.GRPC, new HashMap<>());
          clusterClient.serverListFactory(new ServerListFactory() {
              @Override
              public String genNextServer() {
                  return "127.0.0.1:" + port;
              }
              
              @Override
              public String getCurrentServer() {
                  return "127.0.0.1:" + port;
              }
              
              @Override
              public List<String> getServerList() {
                  return Collections.singletonList("127.0.0.1:" + port);
              }
          });
          // connect to server
          clusterClient.start();
      }
      
      @Test
      void testClientDiscovery() throws NacosException {
          // client judge ability
          assertEquals(AbilityStatus.SUPPORTED, client.getConnectionAbility(AbilityKey.SERVER_FUZZY_WATCH));
          assertEquals(AbilityStatus.NOT_SUPPORTED, client.getConnectionAbility(AbilityKey.SERVER_DISTRIBUTED_LOCK));
      }
      
      @Test
      void testServerDiscoveryAndJudge() throws Exception {
          Map<String, RequestHandler> handlers = (Map<String, RequestHandler>) registryHandlerFields.get(requestHandlerRegistry);
          
          // set handler
          RequestHandler oldRequestHandler = handlers.remove(ConfigQueryRequest.class.getSimpleName());
          handlers.put(ConfigQueryRequest.class.getSimpleName(), new ClientRequestHandler(filters));
          configService.getConfig("test", "DEFAULT_GROUP", 2000);
          // wait server invoke
          Thread.sleep(3000);
          assertTrue(serverSuccess);
          // recover
          handlers.remove(ConfigQueryRequest.class.getSimpleName());
          handlers.put(ConfigQueryRequest.class.getSimpleName(), oldRequestHandler);
      }
      
      @Test
      void testClientJudge() throws Exception {
          List<ServerRequestHandler> handlers = (List<ServerRequestHandler>) serverReuqestHandlersField.get(client);
          handlers.clear();
          // register
          client.registerServerRequestHandler(new ServerRequestHandler() {
              @Override
              public Response requestReply(Request request, Connection connection) {
                  if (connection.getConnectionAbility(AbilityKey.SERVER_FUZZY_WATCH).equals(AbilityStatus.SUPPORTED)
                          && connection.getConnectionAbility(AbilityKey.SERVER_DISTRIBUTED_LOCK).equals(AbilityStatus.NOT_SUPPORTED)) {
                      clientSuccess = true;
                  }
                  return new Response() {
                  };
              }
          });
          
          // get id
          Connection conn = (Connection) currentConnField.get(client);
          
          com.alibaba.nacos.core.remote.Connection connection = connectionManager.getConnection(conn.getConnectionId());
          try {
              connection.request(new NotifySubscriberRequest(), 2000L);
          } catch (NacosException e) {
              // nothing to do
          }
          
          // wait client react
          Thread.sleep(4000);
          assertTrue(clientSuccess);
      }
      
      @Test
      void testClusterClient() throws IllegalAccessException, NacosException, InterruptedException, NoSuchFieldException {
          Map<String, RequestHandler> handlers = (Map<String, RequestHandler>) registryHandlerFields.get(requestHandlerRegistry);
          
          // set handler
          RequestHandler oldRequestHandler = handlers.remove(ConfigQueryRequest.class.getSimpleName());
          handlers.put(ConfigQueryRequest.class.getSimpleName(), new ClusterClientRequestHandler(filters));
          configService.getConfig("test", "DEFAULT_GROUP", 2000);
          // wait server invoke
          Thread.sleep(3000);
          assertTrue(clusterSuccess);
          // recover
          handlers.remove(ConfigQueryRequest.class.getSimpleName());
          handlers.put(ConfigQueryRequest.class.getSimpleName(), oldRequestHandler);
      }
      
      @Test
      void testNegotiationTimeout() throws Exception {
          Object origin = setupRequestHandlerField.get(client);
          // set null for setupRequestHandlerField
          setupRequestHandlerField.set(client, null);
          // try connect
          Connection connection = client.connectToServer(new RpcClient.ServerInfo("127.0.0.1", port));
          assertNull(connection);
          // recovery
          setupRequestHandlerField.set(client, origin);
      }
      
      @AfterEach
      void recover() throws IllegalAccessException, NacosException {
          abstractAbilityControlManager.set(NacosAbilityManagerHolder.class, oldInstance);
          client.shutdown();
      }
      
      /**
       * just to test ability
       */
      class ClientRequestHandler extends RequestHandler<ConfigQueryRequest, ConfigQueryResponse> {
          
          public ClientRequestHandler(RequestFilters requestFilters) throws NoSuchFieldException, IllegalAccessException {
              Field declaredField = RequestHandler.class.getDeclaredField("requestFilters");
              declaredField.setAccessible(true);
              declaredField.set(this, requestFilters);
          }
          
          @Override
          public ConfigQueryResponse handle(ConfigQueryRequest request, RequestMeta meta) throws NacosException {
              if (meta.getConnectionAbility(AbilityKey.SDK_CLIENT_FUZZY_WATCH).equals(AbilityStatus.SUPPORTED)) {
                  serverSuccess = true;
              }
              return new ConfigQueryResponse();
          }
      }
      
      /**
       * just to test ability.
       */
      class ClusterClientRequestHandler extends RequestHandler<ConfigQueryRequest, ConfigQueryResponse> {
          
          public ClusterClientRequestHandler(RequestFilters requestFilters) throws NoSuchFieldException, IllegalAccessException {
              Field declaredField = RequestHandler.class.getDeclaredField("requestFilters");
              declaredField.setAccessible(true);
              declaredField.set(this, requestFilters);
          }
          
          @Override
          public ConfigQueryResponse handle(ConfigQueryRequest request, RequestMeta meta) throws NacosException {
              if (meta.getConnectionAbility(AbilityKey.CLUSTER_CLIENT_TEST_1).equals(AbilityStatus.SUPPORTED)) {
                  clusterSuccess = true;
              }
              return new ConfigQueryResponse();
          }
      }
  }
  