...48import org.openqa.selenium.grid.web.CombinedHandler;49import org.openqa.selenium.grid.web.CommandHandler;50import org.openqa.selenium.net.PortProber;51import org.openqa.selenium.remote.Dialect;52import org.openqa.selenium.remote.NewSessionPayload;53import org.openqa.selenium.remote.SessionId;54import org.openqa.selenium.remote.http.HttpClient;55import org.openqa.selenium.remote.http.HttpRequest;56import org.openqa.selenium.remote.http.HttpResponse;57import org.openqa.selenium.remote.tracing.DistributedTracer;58import org.openqa.selenium.support.ui.FluentWait;59import org.openqa.selenium.support.ui.Wait;60import java.io.IOException;61import java.io.UncheckedIOException;62import java.net.MalformedURLException;63import java.net.URI;64import java.net.URISyntaxException;65import java.net.URL;66import java.time.Duration;67import java.util.Map;68import java.util.UUID;69import java.util.concurrent.atomic.AtomicBoolean;70public class DistributorTest {71 private DistributedTracer tracer;72 private EventBus bus;73 private HttpClient.Factory clientFactory;74 private Distributor local;75 private Distributor distributor;76 private ImmutableCapabilities caps;77 @Before78 public void setUp() throws MalformedURLException {79 tracer = DistributedTracer.builder().build();80 bus = new GuavaEventBus();81 clientFactory = HttpClient.Factory.createDefault();82 LocalSessionMap sessions = new LocalSessionMap(tracer, bus);83 local = new LocalDistributor(tracer, bus, HttpClient.Factory.createDefault(), sessions);84 distributor = new RemoteDistributor(85 tracer,86 new PassthroughHttpClient.Factory<>(local),87 new URL("http://does.not.exist/"));88 caps = new ImmutableCapabilities("browserName", "cheese");89 }90 @Test91 public void creatingANewSessionWithoutANodeEndsInFailure() {92 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {93 assertThatExceptionOfType(SessionNotCreatedException.class)94 .isThrownBy(() -> distributor.newSession(createRequest(payload)));95 }96 }97 @Test98 public void shouldBeAbleToAddANodeAndCreateASession() throws URISyntaxException {99 URI nodeUri = new URI("http://example:5678");100 URI routableUri = new URI("http://localhost:1234");101 LocalSessionMap sessions = new LocalSessionMap(tracer, bus);102 LocalNode node = LocalNode.builder(tracer, bus, clientFactory, routableUri)103 .add(caps, new TestSessionFactory((id, c) -> new Session(id, nodeUri, c)))104 .build();105 Distributor distributor = new LocalDistributor(106 tracer,107 bus,108 new PassthroughHttpClient.Factory<>(node),109 sessions);110 distributor.add(node);111 MutableCapabilities sessionCaps = new MutableCapabilities(caps);112 sessionCaps.setCapability("sausages", "gravy");113 try (NewSessionPayload payload = NewSessionPayload.create(sessionCaps)) {114 Session session = distributor.newSession(createRequest(payload)).getSession();115 assertThat(session.getCapabilities()).isEqualTo(sessionCaps);116 assertThat(session.getUri()).isEqualTo(routableUri);117 }118 }119 @Test120 public void creatingASessionAddsItToTheSessionMap() throws URISyntaxException {121 URI nodeUri = new URI("http://example:5678");122 URI routableUri = new URI("http://localhost:1234");123 LocalSessionMap sessions = new LocalSessionMap(tracer, bus);124 LocalNode node = LocalNode.builder(tracer, bus, clientFactory, routableUri)125 .add(caps, new TestSessionFactory((id, c) -> new Session(id, nodeUri, c)))126 .build();127 Distributor distributor = new LocalDistributor(128 tracer,129 bus,130 new PassthroughHttpClient.Factory<>(node),131 sessions);132 distributor.add(node);133 MutableCapabilities sessionCaps = new MutableCapabilities(caps);134 sessionCaps.setCapability("sausages", "gravy");135 try (NewSessionPayload payload = NewSessionPayload.create(sessionCaps)) {136 Session returned = distributor.newSession(createRequest(payload)).getSession();137 Session session = sessions.get(returned.getId());138 assertThat(session.getCapabilities()).isEqualTo(sessionCaps);139 assertThat(session.getUri()).isEqualTo(routableUri);140 }141 }142 @Test143 public void shouldBeAbleToRemoveANode() throws URISyntaxException, MalformedURLException {144 URI nodeUri = new URI("http://example:5678");145 URI routableUri = new URI("http://localhost:1234");146 LocalSessionMap sessions = new LocalSessionMap(tracer, bus);147 LocalNode node = LocalNode.builder(tracer, bus, clientFactory, routableUri)148 .add(caps, new TestSessionFactory((id, c) -> new Session(id, nodeUri, c)))149 .build();150 Distributor local = new LocalDistributor(151 tracer,152 bus,153 new PassthroughHttpClient.Factory<>(node),154 sessions);155 distributor = new RemoteDistributor(156 tracer,157 new PassthroughHttpClient.Factory<>(local),158 new URL("http://does.not.exist"));159 distributor.add(node);160 distributor.remove(node.getId());161 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {162 assertThatExceptionOfType(SessionNotCreatedException.class)163 .isThrownBy(() -> distributor.newSession(createRequest(payload)));164 }165 }166 @Test167 public void registeringTheSameNodeMultipleTimesOnlyCountsTheFirstTime()168 throws URISyntaxException {169 URI nodeUri = new URI("http://example:5678");170 URI routableUri = new URI("http://localhost:1234");171 LocalNode node = LocalNode.builder(tracer, bus, clientFactory, routableUri)172 .add(caps, new TestSessionFactory((id, c) -> new Session(id, nodeUri, c)))173 .build();174 local.add(node);175 local.add(node);176 DistributorStatus status = local.getStatus();177 assertThat(status.getNodes().size()).isEqualTo(1);178 }179 @Test180 public void theMostLightlyLoadedNodeIsSelectedFirst() {181 // Create enough hosts so that we avoid the scheduler returning hosts in:182 // * insertion order183 // * reverse insertion order184 // * sorted with most heavily used first185 SessionMap sessions = new LocalSessionMap(tracer, bus);186 Node lightest = createNode(caps, 10, 0);187 Node medium = createNode(caps, 10, 4);188 Node heavy = createNode(caps, 10, 6);189 Node massive = createNode(caps, 10, 8);190 CombinedHandler handler = new CombinedHandler();191 handler.addHandler(lightest);192 handler.addHandler(medium);193 handler.addHandler(heavy);194 handler.addHandler(massive);195 Distributor distributor = new LocalDistributor(196 tracer,197 bus,198 new PassthroughHttpClient.Factory<>(handler),199 sessions)200 .add(heavy)201 .add(medium)202 .add(lightest)203 .add(massive);204 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {205 Session session = distributor.newSession(createRequest(payload)).getSession();206 assertThat(session.getUri()).isEqualTo(lightest.getStatus().getUri());207 }208 }209 @Test210 public void shouldUseLastSessionCreatedTimeAsTieBreaker() {211 SessionMap sessions = new LocalSessionMap(tracer, bus);212 Node leastRecent = createNode(caps, 5, 0);213 CombinedHandler handler = new CombinedHandler();214 handler.addHandler(sessions);215 handler.addHandler(leastRecent);216 Distributor distributor = new LocalDistributor(217 tracer,218 bus,219 new PassthroughHttpClient.Factory<>(handler),220 sessions)221 .add(leastRecent);222 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {223 distributor.newSession(createRequest(payload));224 // Will be "leastRecent" by default225 }226 Node middle = createNode(caps, 5, 0);227 handler.addHandler(middle);228 distributor.add(middle);229 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {230 Session session = distributor.newSession(createRequest(payload)).getSession();231 // Least lightly loaded is middle232 assertThat(session.getUri()).isEqualTo(middle.getStatus().getUri());233 }234 Node mostRecent = createNode(caps, 5, 0);235 handler.addHandler(mostRecent);236 distributor.add(mostRecent);237 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {238 Session session = distributor.newSession(createRequest(payload)).getSession();239 // Least lightly loaded is most recent240 assertThat(session.getUri()).isEqualTo(mostRecent.getStatus().getUri());241 }242 // All the nodes should be equally loaded.243 Map<Capabilities, Integer> expected = mostRecent.getStatus().getStereotypes();244 assertThat(leastRecent.getStatus().getStereotypes()).isEqualTo(expected);245 assertThat(middle.getStatus().getStereotypes()).isEqualTo(expected);246 // All nodes are now equally loaded. We should be going in time order now247 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {248 Session session = distributor.newSession(createRequest(payload)).getSession();249 assertThat(session.getUri()).isEqualTo(leastRecent.getStatus().getUri());250 }251 }252 @Test253 public void shouldIncludeHostsThatAreUpInHostList() {254 CombinedHandler handler = new CombinedHandler();255 SessionMap sessions = new LocalSessionMap(tracer, bus);256 handler.addHandler(sessions);257 URI uri = createUri();258 Node alwaysDown = LocalNode.builder(tracer, bus, clientFactory, uri)259 .add(caps, new TestSessionFactory((id, c) -> new Session(id, uri, c)))260 .advanced()261 .healthCheck(() -> new HealthCheck.Result(false, "Boo!"))262 .build();263 handler.addHandler(alwaysDown);264 Node alwaysUp = LocalNode.builder(tracer, bus, clientFactory, uri)265 .add(caps, new TestSessionFactory((id, c) -> new Session(id, uri, c)))266 .advanced()267 .healthCheck(() -> new HealthCheck.Result(true, "Yay!"))268 .build();269 handler.addHandler(alwaysUp);270 LocalDistributor distributor = new LocalDistributor(271 tracer,272 bus,273 new PassthroughHttpClient.Factory<>(handler),274 sessions);275 handler.addHandler(distributor);276 distributor.add(alwaysDown);277 // Should be unable to create a session because the node is down.278 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {279 assertThatExceptionOfType(SessionNotCreatedException.class)280 .isThrownBy(() -> distributor.newSession(createRequest(payload)));281 }282 distributor.add(alwaysUp);283 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {284 distributor.newSession(createRequest(payload));285 }286 }287 @Test288 public void shouldNotScheduleAJobIfAllSlotsAreBeingUsed() {289 SessionMap sessions = new LocalSessionMap(tracer, bus);290 CombinedHandler handler = new CombinedHandler();291 Distributor distributor = new LocalDistributor(292 tracer,293 bus,294 new PassthroughHttpClient.Factory<>(handler),295 sessions);296 handler.addHandler(distributor);297 Node node = createNode(caps, 1, 0);298 handler.addHandler(node);299 distributor.add(node);300 // Use up the one slot available301 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {302 distributor.newSession(createRequest(payload));303 }304 // Now try and create a session.305 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {306 assertThatExceptionOfType(SessionNotCreatedException.class)307 .isThrownBy(() -> distributor.newSession(createRequest(payload)));308 }309 }310 @Test311 public void shouldReleaseSlotOnceSessionEnds() {312 SessionMap sessions = new LocalSessionMap(tracer, bus);313 CombinedHandler handler = new CombinedHandler();314 Distributor distributor = new LocalDistributor(315 tracer,316 bus,317 new PassthroughHttpClient.Factory<>(handler),318 sessions);319 handler.addHandler(distributor);320 Node node = createNode(caps, 1, 0);321 handler.addHandler(node);322 distributor.add(node);323 // Use up the one slot available324 Session session;325 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {326 session = distributor.newSession(createRequest(payload)).getSession();327 }328 // Make sure the session map has the session329 sessions.get(session.getId());330 node.stop(session.getId());331 // Now wait for the session map to say the session is gone.332 Wait<Object> wait = new FluentWait<>(new Object()).withTimeout(Duration.ofSeconds(2));333 wait.until(obj -> {334 try {335 sessions.get(session.getId());336 return false;337 } catch (NoSuchSessionException e) {338 return true;339 }340 });341 wait.until(obj -> distributor.getStatus().hasCapacity());342 // And we should now be able to create another session.343 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {344 distributor.newSession(createRequest(payload));345 }346 }347 @Test348 public void shouldNotStartASessionIfTheCapabilitiesAreNotSupported() {349 CombinedHandler handler = new CombinedHandler();350 LocalSessionMap sessions = new LocalSessionMap(tracer, bus);351 handler.addHandler(handler);352 Distributor distributor = new LocalDistributor(353 tracer,354 bus,355 new PassthroughHttpClient.Factory<>(handler),356 sessions);357 handler.addHandler(distributor);358 Node node = createNode(caps, 1, 0);359 handler.addHandler(node);360 distributor.add(node);361 ImmutableCapabilities unmatched = new ImmutableCapabilities("browserName", "transit of venus");362 try (NewSessionPayload payload = NewSessionPayload.create(unmatched)) {363 assertThatExceptionOfType(SessionNotCreatedException.class)364 .isThrownBy(() -> distributor.newSession(createRequest(payload)));365 }366 }367 @Test368 public void attemptingToStartASessionWhichFailsMarksAsTheSlotAsAvailable() {369 CombinedHandler handler = new CombinedHandler();370 SessionMap sessions = new LocalSessionMap(tracer, bus);371 handler.addHandler(sessions);372 URI uri = createUri();373 Node node = LocalNode.builder(tracer, bus, clientFactory, uri)374 .add(caps, new TestSessionFactory((id, caps) -> {375 throw new SessionNotCreatedException("OMG");376 }))377 .build();378 handler.addHandler(node);379 Distributor distributor = new LocalDistributor(380 tracer,381 bus,382 new PassthroughHttpClient.Factory<>(handler),383 sessions);384 handler.addHandler(distributor);385 distributor.add(node);386 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {387 assertThatExceptionOfType(SessionNotCreatedException.class)388 .isThrownBy(() -> distributor.newSession(createRequest(payload)));389 }390 assertThat(distributor.getStatus().hasCapacity()).isTrue();391 }392 @Test393 public void shouldReturnNodesThatWereDownToPoolOfNodesOnceTheyMarkTheirHealthCheckPasses() {394 CombinedHandler handler = new CombinedHandler();395 SessionMap sessions = new LocalSessionMap(tracer, bus);396 handler.addHandler(sessions);397 AtomicBoolean isUp = new AtomicBoolean(false);398 URI uri = createUri();399 Node node = LocalNode.builder(tracer, bus, clientFactory, uri)400 .add(caps, new TestSessionFactory((id, caps) -> new Session(id, uri, caps)))401 .advanced()402 .healthCheck(() -> new HealthCheck.Result(isUp.get(), "TL;DR"))403 .build();404 handler.addHandler(node);405 LocalDistributor distributor = new LocalDistributor(406 tracer,407 bus,408 new PassthroughHttpClient.Factory<>(handler),409 sessions);410 handler.addHandler(distributor);411 distributor.add(node);412 // Should be unable to create a session because the node is down.413 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {414 assertThatExceptionOfType(SessionNotCreatedException.class)415 .isThrownBy(() -> distributor.newSession(createRequest(payload)));416 }417 // Mark the node as being up418 isUp.set(true);419 // Kick the machinery to ensure that everything is fine.420 distributor.refresh();421 // Because the node is now up and running, we should now be able to create a session422 try (NewSessionPayload payload = NewSessionPayload.create(caps)) {423 distributor.newSession(createRequest(payload));424 }425 }426 @Test427 @Ignore428 public void shouldPriotizeHostsWithTheMostSlotsAvailableForASessionType() {429 // Consider the case where you have 1 Windows machine and 5 linux machines. All of these hosts430 // can run Chrome and Firefox sessions, but only one can run Edge sessions. Ideally, the machine431 // able to run Edge would be sorted last.432 fail("Write me");433 }434 private Node createNode(Capabilities stereotype, int count, int currentLoad) {435 URI uri = createUri();436 LocalNode.Builder builder = LocalNode.builder(tracer, bus, clientFactory, uri);437 for (int i = 0; i < count; i++) {438 builder.add(stereotype, new TestSessionFactory((id, caps) -> new HandledSession(uri, caps)));439 }440 LocalNode node = builder.build();441 for (int i = 0; i < currentLoad; i++) {442 // Ignore the session. We're just creating load.443 node.newSession(new CreateSessionRequest(444 ImmutableSet.copyOf(Dialect.values()),445 stereotype,446 ImmutableMap.of()));447 }448 return node;449 }450 @Test451 @Ignore452 public void shouldCorrectlySetSessionCountsWhenStartedAfterNodeWithSession() {453 fail("write me");454 }455 @Test456 public void statusShouldIndicateThatDistributorIsNotAvailableIfNodesAreDown()457 throws URISyntaxException {458 Capabilities capabilities = new ImmutableCapabilities("cheese", "peas");459 URI uri = new URI("http://exmaple.com");460 Node node = LocalNode.builder(tracer, bus, clientFactory, uri)461 .add(capabilities, new TestSessionFactory((id, caps) -> new Session(id, uri, caps)))462 .advanced()463 .healthCheck(() -> new HealthCheck.Result(false, "TL;DR"))464 .build();465 local.add(node);466 DistributorStatus status = local.getStatus();467 assertFalse(status.hasCapacity());468 }469 private HttpRequest createRequest(NewSessionPayload payload) {470 StringBuilder builder = new StringBuilder();471 try {472 payload.writeTo(builder);473 } catch (IOException e) {474 throw new UncheckedIOException(e);475 }476 HttpRequest request = new HttpRequest(POST, "/se/grid/distributor/session");477 request.setContent(utf8String(builder.toString()));478 return request;479 }480 private URI createUri() {481 try {482 return new URI("http://localhost:" + PortProber.findFreePort());483 } catch (URISyntaxException e) {...