code: package org.openqa.selenium.grid.distributor;
  
  import static org.assertj.core.api.Assertions.assertThat;
  import static org.openqa.selenium.grid.data.Availability.DOWN;
  import static org.openqa.selenium.grid.data.Availability.UP;
  
  import java.net.MalformedURLException;
  import java.net.URI;
  import java.net.URL;
  import java.time.Duration;
  import java.time.Instant;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.concurrent.atomic.AtomicReference;
  import org.junit.jupiter.api.Test;
  import org.openqa.selenium.MutableCapabilities;
  import org.openqa.selenium.SessionNotCreatedException;
  import org.openqa.selenium.grid.data.Availability;
  import org.openqa.selenium.grid.data.CreateSessionResponse;
  import org.openqa.selenium.grid.data.DefaultSlotMatcher;
  import org.openqa.selenium.grid.data.DistributorStatus;
  import org.openqa.selenium.grid.data.Session;
  import org.openqa.selenium.grid.distributor.local.LocalDistributor;
  import org.openqa.selenium.grid.distributor.remote.RemoteDistributor;
  import org.openqa.selenium.grid.distributor.selector.DefaultSlotSelector;
  import org.openqa.selenium.grid.node.HealthCheck;
  import org.openqa.selenium.grid.node.Node;
  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.grid.web.CombinedHandler;
  import org.openqa.selenium.internal.Either;
  
  public class DistributorNodeAvailabilityTest extends DistributorTestBase {
  
    @Test
    void registeringTheSameNodeMultipleTimesOnlyCountsTheFirstTime() {
      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.add(node);
  
      DistributorStatus status = local.getStatus();
  
      assertThat(status.getNodes().size()).isEqualTo(1);
    }
  
    @Test
    void shouldBeAbleToRemoveANode() throws MalformedURLException {
      LocalSessionMap 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);
      Distributor distributor =
          new RemoteDistributor(
              tracer,
              new PassthroughHttpClient.Factory(local),
              new URL("http://does.not.exist"),
              registrationSecret);
      distributor.add(node);
      distributor.remove(node.getId());
  
      Either<SessionNotCreatedException, CreateSessionResponse> result =
          local.newSession(createRequest(caps));
      assertThatEither(result).isLeft();
    }
  
    @Test
    void shouldIncludeHostsThatAreUpInHostList() {
      CombinedHandler handler = new CombinedHandler();
  
      SessionMap sessions = new LocalSessionMap(tracer, bus);
      NewSessionQueue queue =
          new LocalNewSessionQueue(
              tracer,
              new DefaultSlotMatcher(),
              Duration.ofSeconds(2),
              Duration.ofSeconds(2),
              registrationSecret,
              5);
      handler.addHandler(sessions);
  
      URI uri = createUri();
      Node alwaysDown =
          LocalNode.builder(tracer, bus, uri, uri, registrationSecret)
              .add(
                  caps,
                  new TestSessionFactory(
                      (id, c) -> new Session(id, uri, stereotype, c, Instant.now())))
              .advanced()
              .healthCheck(() -> new HealthCheck.Result(DOWN, "Boo!"))
              .build();
      handler.addHandler(alwaysDown);
  
      local =
          new LocalDistributor(
              tracer,
              bus,
              new PassthroughHttpClient.Factory(handler),
              sessions,
              queue,
              new DefaultSlotSelector(),
              registrationSecret,
              Duration.ofSeconds(1),
              false,
              Duration.ofSeconds(5),
              newSessionThreadPoolSize);
      handler.addHandler(local);
      local.add(alwaysDown);
      waitForAllNodesToMeetCondition(local, 1, DOWN);
  
      // Should be unable to create a session because the node is down.
      Either<SessionNotCreatedException, CreateSessionResponse> result =
          local.newSession(createRequest(caps));
      assertThatEither(result).isLeft();
  
      Node alwaysUp =
          LocalNode.builder(tracer, bus, uri, uri, registrationSecret)
              .add(
                  caps,
                  new TestSessionFactory(
                      (id, c) -> new Session(id, uri, stereotype, c, Instant.now())))
              .advanced()
              .healthCheck(() -> new HealthCheck.Result(UP, "Yay!"))
              .build();
      handler.addHandler(alwaysUp);
  
      local.add(alwaysUp);
      waitToHaveCapacity(local);
  
      result = local.newSession(createRequest(caps));
      assertThatEither(result).isRight();
    }
  
    @Test
    void shouldNotRemoveNodeWhoseHealthCheckPassesBeforeThreshold() throws InterruptedException {
      CombinedHandler handler = new CombinedHandler();
  
      AtomicInteger count = new AtomicInteger(0);
      CountDownLatch latch = new CountDownLatch(1);
  
      SessionMap sessions = new LocalSessionMap(tracer, bus);
      handler.addHandler(sessions);
      NewSessionQueue queue =
          new LocalNewSessionQueue(
              tracer,
              new DefaultSlotMatcher(),
              Duration.ofSeconds(2),
              Duration.ofSeconds(2),
              registrationSecret,
              5);
  
      URI uri = createUri();
      Node node =
          LocalNode.builder(tracer, bus, uri, uri, registrationSecret)
              .add(
                  caps,
                  new TestSessionFactory(
                      (id, caps) -> new Session(id, uri, stereotype, caps, Instant.now())))
              .advanced()
              .healthCheck(
                  () -> {
                    if (count.get() <= 4) {
                      count.incrementAndGet();
                      return new HealthCheck.Result(DOWN, "Down");
                    }
                    latch.countDown();
                    return new HealthCheck.Result(UP, "Up");
                  })
              .build();
      handler.addHandler(node);
  
      local =
          new LocalDistributor(
              tracer,
              bus,
              new PassthroughHttpClient.Factory(handler),
              sessions,
              queue,
              new DefaultSlotSelector(),
              registrationSecret,
              Duration.ofSeconds(1),
              false,
              Duration.ofSeconds(5),
              newSessionThreadPoolSize);
      handler.addHandler(local);
      local.add(node);
  
      latch.await(60, TimeUnit.SECONDS);
  
      waitToHaveCapacity(local);
  
      Either<SessionNotCreatedException, CreateSessionResponse> result =
          local.newSession(createRequest(caps));
      assertThatEither(result).isRight();
    }
  
    @Test
    void shouldReturnNodesThatWereDownToPoolOfNodesOnceTheyMarkTheirHealthCheckPasses() {
      CombinedHandler handler = new CombinedHandler();
  
      SessionMap sessions = new LocalSessionMap(tracer, bus);
      handler.addHandler(sessions);
      AtomicReference<Availability> isUp = new AtomicReference<>(DOWN);
      NewSessionQueue queue =
          new LocalNewSessionQueue(
              tracer,
              new DefaultSlotMatcher(),
              Duration.ofSeconds(2),
              Duration.ofSeconds(2),
              registrationSecret,
              5);
  
      URI uri = createUri();
      Node node =
          LocalNode.builder(tracer, bus, uri, uri, registrationSecret)
              .add(
                  caps,
                  new TestSessionFactory(
                      (id, caps) -> new Session(id, uri, stereotype, caps, Instant.now())))
              .advanced()
              .healthCheck(() -> new HealthCheck.Result(isUp.get(), "TL;DR"))
              .build();
      handler.addHandler(node);
  
      local =
          new LocalDistributor(
              tracer,
              bus,
              new PassthroughHttpClient.Factory(handler),
              sessions,
              queue,
              new DefaultSlotSelector(),
              registrationSecret,
              Duration.ofSeconds(1),
              false,
              Duration.ofSeconds(5),
              newSessionThreadPoolSize);
      handler.addHandler(local);
      local.add(node);
      waitForAllNodesToMeetCondition(local, 1, DOWN);
  
      // Should be unable to create a session because the node is down.
      Either<SessionNotCreatedException, CreateSessionResponse> result =
          local.newSession(createRequest(caps));
      assertThatEither(result).isLeft();
  
      // Mark the node as being up
      isUp.set(UP);
      // Kick the machinery to ensure that everything is fine.
      local.refresh();
  
      // Because the node is now up and running, we should now be able to create a session
      result = local.newSession(createRequest(caps));
      assertThatEither(result).isRight();
    }
  
    @Test
    void shouldBeAbleToAddANodeAndCreateASession() {
      LocalSessionMap 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);
      waitToHaveCapacity(local);
  
      MutableCapabilities sessionCaps = new MutableCapabilities(caps);
      sessionCaps.setCapability("sausages", "gravy");
  
      Either<SessionNotCreatedException, CreateSessionResponse> result =
          local.newSession(createRequest(sessionCaps));
      assertThatEither(result).isRight();
      Session session = result.right().getSession();
      assertThat(session.getCapabilities().getCapability("sausages"))
          .isEqualTo(sessionCaps.getCapability("sausages"));
      assertThat(session.getUri()).isEqualTo(routableUri);
    }
  }
  