@Test
    public void testLargeLocks() {
        var tenantId = new TenantId("tenant");
        var applicationInstanceId = new ApplicationInstanceId("app:dev:us-east-1:default");
        var applicationInstanceReference = new ApplicationInstanceReference(tenantId, applicationInstanceId);

        var policy = mock(HostedVespaPolicy.class);
        var zookeeperStatusService = mock(ZkStatusService.class);
        var serviceMonitor = mock(ServiceMonitor.class);
        var applicationInstance = mock(ApplicationInstance.class);
        var clusterControllerClientFactory = mock(ClusterControllerClientFactory.class);
        var clock = new ManualClock();
        var applicationApiFactory = mock(ApplicationApiFactory.class);
        var lock = mock(ApplicationLock.class);

        when(policy.grantSuspensionRequest(any(), any())).thenReturn(SuspensionReasons.nothingNoteworthy());
        when(serviceMonitor.getApplication(any(HostName.class))).thenReturn(Optional.of(applicationInstance));
        when(applicationInstance.reference()).thenReturn(applicationInstanceReference);
        when(zookeeperStatusService.lockApplication(any(), any())).thenReturn(lock);
        when(lock.getApplicationInstanceStatus()).thenReturn(NO_REMARKS);

        var orchestrator = new OrchestratorImpl(
                policy,
                clusterControllerClientFactory,
                zookeeperStatusService,
                serviceMonitor,
                20,
                clock,
                applicationApiFactory,
                flagSource);

        HostName parentHostname = new HostName("parent.vespa.ai");

        verify(serviceMonitor, atLeastOnce()).registerListener(zookeeperStatusService);
        verifyNoMoreInteractions(serviceMonitor);

        orchestrator.suspendAll(parentHostname, List.of(parentHostname));

        ArgumentCaptor<OrchestratorContext> contextCaptor = ArgumentCaptor.forClass(OrchestratorContext.class);
        verify(zookeeperStatusService, times(2)).lockApplication(contextCaptor.capture(), any());
        List<OrchestratorContext> contexts = contextCaptor.getAllValues();

        // First invocation is probe, second is not.
        assertEquals(2, contexts.size());
        assertTrue(contexts.get(0).isProbe());
        assertTrue(contexts.get(0).largeLocks());
        assertFalse(contexts.get(1).isProbe());
        assertTrue(contexts.get(1).largeLocks());

        verify(applicationApiFactory, times(2)).create(any(), any(), any());
        verify(policy, times(2)).grantSuspensionRequest(any(), any());
        verify(serviceMonitor, atLeastOnce()).getApplication(any(HostName.class));
        verify(lock, times(2)).getApplicationInstanceStatus();

        // Each zookeeperStatusService that is created, is closed.
        verify(zookeeperStatusService, times(2)).lockApplication(any(), any());
        verify(lock, times(2)).close();

        verifyNoMoreInteractions(
                policy,
                clusterControllerClientFactory,
                zookeeperStatusService,
                lock,
                serviceMonitor,
                applicationApiFactory);
    }
Test Orchestrator with a mock backend (the MockCurator)

@author smorgrav
orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java