code: package org.openqa.selenium.grid.distributor;
  
  import static org.assertj.core.api.Assertions.assertThat;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  import java.time.Duration;
  import java.time.Instant;
  import java.util.Set;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.TimeUnit;
  import org.junit.jupiter.api.Test;
  import org.openqa.selenium.SessionNotCreatedException;
  import org.openqa.selenium.grid.data.CreateSessionResponse;
  import org.openqa.selenium.grid.data.DefaultSlotMatcher;
  import org.openqa.selenium.grid.data.NodeDrainComplete;
  import org.openqa.selenium.grid.data.NodeStatus;
  import org.openqa.selenium.grid.data.Session;
  import org.openqa.selenium.grid.distributor.local.LocalDistributor;
  import org.openqa.selenium.grid.distributor.selector.DefaultSlotSelector;
  import org.openqa.selenium.grid.node.local.LocalNode;
  import org.openqa.selenium.grid.sessionmap.SessionMap;
  import org.openqa.selenium.grid.sessionmap.local.LocalSessionMap;
  import org.openqa.selenium.grid.sessionqueue.NewSessionQueue;
  import org.openqa.selenium.grid.sessionqueue.local.LocalNewSessionQueue;
  import org.openqa.selenium.grid.testing.PassthroughHttpClient;
  import org.openqa.selenium.grid.testing.TestSessionFactory;
  import org.openqa.selenium.internal.Either;
  import org.openqa.selenium.support.ui.FluentWait;
  
  public class DistributorDrainingTest extends DistributorTestBase {
    @Test
    void drainedNodeDoesNotShutDownIfNotEmpty() throws InterruptedException {
      SessionMap sessions = new LocalSessionMap(tracer, bus);
      NewSessionQueue queue =
          new LocalNewSessionQueue(
              tracer,
              new DefaultSlotMatcher(),
              Duration.ofSeconds(2),
              Duration.ofSeconds(2),
              registrationSecret,
              5);
      LocalNode node =
          LocalNode.builder(tracer, bus, routableUri, routableUri, registrationSecret)
              .add(
                  caps,
                  new TestSessionFactory(
                      (id, c) -> new Session(id, nodeUri, stereotype, c, Instant.now())))
              .build();
  
      CountDownLatch latch = new CountDownLatch(1);
      bus.addListener(NodeDrainComplete.listener(ignored -> latch.countDown()));
  
      local =
          new LocalDistributor(
              tracer,
              bus,
              new PassthroughHttpClient.Factory(node),
              sessions,
              queue,
              new DefaultSlotSelector(),
              registrationSecret,
              Duration.ofMinutes(5),
              false,
              Duration.ofSeconds(5),
              newSessionThreadPoolSize);
      local.add(node);
      waitToHaveCapacity(local);
  
      Either<SessionNotCreatedException, CreateSessionResponse> session =
          local.newSession(createRequest(caps));
      assertThatEither(session).isRight();
  
      local.drain(node.getId());
  
      latch.await(5, TimeUnit.SECONDS);
  
      assertThat(latch.getCount()).isEqualTo(1);
  
      assertThat(local.getStatus().getNodes().size()).isEqualTo(1);
    }
  
    @Test
    void drainedNodeShutsDownAfterSessionsFinish() throws InterruptedException {
      SessionMap sessions = new LocalSessionMap(tracer, bus);
      NewSessionQueue queue =
          new LocalNewSessionQueue(
              tracer,
              new DefaultSlotMatcher(),
              Duration.ofSeconds(2),
              Duration.ofSeconds(2),
              registrationSecret,
              5);
      LocalNode node =
          LocalNode.builder(tracer, bus, routableUri, routableUri, registrationSecret)
              .add(
                  caps,
                  new TestSessionFactory(
                      (id, c) -> new Session(id, nodeUri, stereotype, c, Instant.now())))
              .add(
                  caps,
                  new TestSessionFactory(
                      (id, c) -> new Session(id, nodeUri, stereotype, c, Instant.now())))
              .maximumConcurrentSessions(5)
              .build();
  
      CountDownLatch latch = new CountDownLatch(1);
      bus.addListener(NodeDrainComplete.listener(ignored -> latch.countDown()));
  
      local =
          new LocalDistributor(
              tracer,
              bus,
              new PassthroughHttpClient.Factory(node),
              sessions,
              queue,
              new DefaultSlotSelector(),
              registrationSecret,
              Duration.ofMinutes(5),
              false,
              Duration.ofSeconds(5),
              newSessionThreadPoolSize);
      local.add(node);
      waitToHaveCapacity(local);
  
      Either<SessionNotCreatedException, CreateSessionResponse> firstResponse =
          local.newSession(createRequest(caps));
  
      Either<SessionNotCreatedException, CreateSessionResponse> secondResponse =
          local.newSession(createRequest(caps));
  
      local.drain(node.getId());
  
      assertThat(local.getStatus().getNodes().size()).isEqualTo(1);
  
      node.stop(firstResponse.right().getSession().getId());
      node.stop(secondResponse.right().getSession().getId());
  
      latch.await(5, TimeUnit.SECONDS);
  
      waitTillNodesAreRemoved(local);
  
      assertThat(latch.getCount()).isZero();
      assertThat(local.getStatus().getNodes()).isEmpty();
    }
  
    private void waitTillNodesAreRemoved(Distributor distributor) {
      new FluentWait<>(distributor)
          .withTimeout(Duration.ofSeconds(60))
          .pollingEvery(Duration.ofMillis(100))
          .until(
              d -> {
                Set<NodeStatus> nodes = d.getStatus().getNodes();
                return nodes.size() == 0;
              });
    }
  
    @Test
    void testDrainedNodeShutsDownOnceEmpty() throws InterruptedException {
      SessionMap sessions = new LocalSessionMap(tracer, bus);
      NewSessionQueue queue =
          new LocalNewSessionQueue(
              tracer,
              new DefaultSlotMatcher(),
              Duration.ofSeconds(2),
              Duration.ofSeconds(2),
              registrationSecret,
              5);
      LocalNode node =
          LocalNode.builder(tracer, bus, routableUri, routableUri, registrationSecret)
              .add(
                  caps,
                  new TestSessionFactory(
                      (id, c) -> new Session(id, nodeUri, stereotype, c, Instant.now())))
              .build();
  
      CountDownLatch latch = new CountDownLatch(1);
      bus.addListener(NodeDrainComplete.listener(ignored -> latch.countDown()));
  
      local =
          new LocalDistributor(
              tracer,
              bus,
              new PassthroughHttpClient.Factory(node),
              sessions,
              queue,
              new DefaultSlotSelector(),
              registrationSecret,
              Duration.ofMinutes(5),
              false,
              Duration.ofSeconds(5),
              newSessionThreadPoolSize);
      local.add(node);
      waitToHaveCapacity(local);
  
      local.drain(node.getId());
  
      latch.await(5, TimeUnit.SECONDS);
  
      assertThat(latch.getCount()).isZero();
  
      assertThat(local.getStatus().getNodes()).isEmpty();
  
      Either<SessionNotCreatedException, CreateSessionResponse> result =
          local.newSession(createRequest(caps));
      assertThatEither(result).isLeft();
    }
  
    @Test
    void drainingNodeDoesNotAcceptNewSessions() {
      SessionMap sessions = new LocalSessionMap(tracer, bus);
      NewSessionQueue queue =
          new LocalNewSessionQueue(
              tracer,
              new DefaultSlotMatcher(),
              Duration.ofSeconds(2),
              Duration.ofSeconds(2),
              registrationSecret,
              5);
      LocalNode node =
          LocalNode.builder(tracer, bus, routableUri, routableUri, registrationSecret)
              .add(
                  caps,
                  new TestSessionFactory(
                      (id, c) -> new Session(id, nodeUri, stereotype, c, Instant.now())))
              .build();
  
      local =
          new LocalDistributor(
              tracer,
              bus,
              new PassthroughHttpClient.Factory(node),
              sessions,
              queue,
              new DefaultSlotSelector(),
              registrationSecret,
              Duration.ofMinutes(5),
              false,
              Duration.ofSeconds(5),
              newSessionThreadPoolSize);
      local.add(node);
      local.drain(node.getId());
  
      assertTrue(node.isDraining());
  
      Either<SessionNotCreatedException, CreateSessionResponse> result =
          local.newSession(createRequest(caps));
      assertThatEither(result).isLeft();
    }
  }
  