<commit_before>// Copyright 2015-2018 The NATS Authors
// 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 io.nats.client.impl;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;

import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.LockSupport;

import org.junit.Test;

import io.nats.client.Connection;
import io.nats.client.Dispatcher;
import io.nats.client.Message;
import io.nats.client.Nats;
import io.nats.client.NatsTestServer;
import io.nats.client.Options;
import io.nats.client.Subscription;
import io.nats.client.TestHandler;
import io.nats.client.ConnectionListener.Events;

public class DrainTests {

    @Test
    public void testSimpleSubDrain() throws Exception {
        try (NatsTestServer ts = new NatsTestServer(false);
                Connection subCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build());
                Connection pubCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build())) {
            assertTrue("Connected Status", Connection.Status.CONNECTED == subCon.getStatus());
            assertTrue("Connected Status", Connection.Status.CONNECTED == pubCon.getStatus());

            Subscription sub = subCon.subscribe("draintest");
            subCon.flush(Duration.ofSeconds(1)); // Get the sub to the server

            pubCon.publish("draintest", null);
            pubCon.publish("draintest", null); // publish 2
            pubCon.flush(Duration.ofSeconds(1));

            Message msg = sub.nextMessage(Duration.ofSeconds(1)); // read 1
            assertNotNull(msg);

            subCon.flush(Duration.ofSeconds(1));
            CompletableFuture<Boolean> tracker = sub.drain(Duration.ofSeconds(1));

            msg = sub.nextMessage(Duration.ofSeconds(1)); // read the second one, should be there because we drained
            assertNotNull(msg);

            assertTrue(tracker.get(1, TimeUnit.SECONDS));
            assertFalse(sub.isActive());
            assertEquals(((NatsConnection) subCon).getConsumerCount(), 0);
        }
    }

    @Test
    public void testSimpleDispatchDrain() throws Exception {
        try (NatsTestServer ts = new NatsTestServer(false);
                Connection subCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build());
                Connection pubCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build())) {
            assertTrue("Connected Status", Connection.Status.CONNECTED == subCon.getStatus());
            assertTrue("Connected Status", Connection.Status.CONNECTED == pubCon.getStatus());

            AtomicInteger count = new AtomicInteger();
            Dispatcher d = subCon.createDispatcher((msg) -> {
                count.incrementAndGet();
                try {
                    Thread.sleep(2000); // go slow so the main app can drain us
                } catch (Exception e) {

                }
            });
            d.subscribe("draintest");
            subCon.flush(Duration.ofSeconds(1)); // Get the sub to the server

            pubCon.publish("draintest", null);
            pubCon.publish("draintest", null);
            pubCon.flush(Duration.ofSeconds(1));
            subCon.flush(Duration.ofSeconds(1));

            // Drain will unsub the dispatcher, only messages that already arrived
            // are there
            CompletableFuture<Boolean> tracker = d.drain(Duration.ofSeconds(4));

            assertTrue(tracker.get(5, TimeUnit.SECONDS)); // wait for the drain to complete
            assertEquals(count.get(), 2); // Should get both
            assertFalse(d.isActive());
            assertEquals(((NatsConnection) subCon).getConsumerCount(), 0);
        }
    }

    @Test
    public void testSimpleConnectionDrain() throws Exception {
        try (NatsTestServer ts = new NatsTestServer(false);
                Connection subCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build());
                Connection pubCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build())) {
            assertTrue("Connected Status", Connection.Status.CONNECTED == subCon.getStatus());
            assertTrue("Connected Status", Connection.Status.CONNECTED == pubCon.getStatus());

            AtomicInteger count = new AtomicInteger();
            Dispatcher d = subCon.createDispatcher((msg) -> {
                count.incrementAndGet();
                try {<commit_msg>Add dispatcher sub with message handler to DrainTests<commit_after>// Copyright 2015-2018 The NATS Authors
// 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 io.nats.client.impl;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;

import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.LockSupport;

import org.junit.Test;

import io.nats.client.Connection;
import io.nats.client.Dispatcher;
import io.nats.client.Message;
import io.nats.client.Nats;
import io.nats.client.NatsTestServer;
import io.nats.client.Options;
import io.nats.client.Subscription;
import io.nats.client.TestHandler;
import io.nats.client.ConnectionListener.Events;

public class DrainTests {

    @Test
    public void testSimpleSubDrain() throws Exception {
        try (NatsTestServer ts = new NatsTestServer(false);
                Connection subCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build());
                Connection pubCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build())) {
            assertTrue("Connected Status", Connection.Status.CONNECTED == subCon.getStatus());
            assertTrue("Connected Status", Connection.Status.CONNECTED == pubCon.getStatus());

            Subscription sub = subCon.subscribe("draintest");
            subCon.flush(Duration.ofSeconds(1)); // Get the sub to the server

            pubCon.publish("draintest", null);
            pubCon.publish("draintest", null); // publish 2
            pubCon.flush(Duration.ofSeconds(1));

            Message msg = sub.nextMessage(Duration.ofSeconds(1)); // read 1
            assertNotNull(msg);

            subCon.flush(Duration.ofSeconds(1));
            CompletableFuture<Boolean> tracker = sub.drain(Duration.ofSeconds(1));

            msg = sub.nextMessage(Duration.ofSeconds(1)); // read the second one, should be there because we drained
            assertNotNull(msg);

            assertTrue(tracker.get(1, TimeUnit.SECONDS));
            assertFalse(sub.isActive());
            assertEquals(((NatsConnection) subCon).getConsumerCount(), 0);
        }
    }

    @Test
    public void testSimpleDispatchDrain() throws Exception {
        try (NatsTestServer ts = new NatsTestServer(false);
                Connection subCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build());
                Connection pubCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build())) {
            assertTrue("Connected Status", Connection.Status.CONNECTED == subCon.getStatus());
            assertTrue("Connected Status", Connection.Status.CONNECTED == pubCon.getStatus());

            AtomicInteger count = new AtomicInteger();
            Dispatcher d = subCon.createDispatcher((msg) -> {
                count.incrementAndGet();
                try {
                    Thread.sleep(2000); // go slow so the main app can drain us
                } catch (Exception e) {

                }
            });
            d.subscribe("draintest");
            d.subscribe("draintest", (msg) -> { count.incrementAndGet(); });
            subCon.flush(Duration.ofSeconds(1)); // Get the sub to the server

            pubCon.publish("draintest", null);
            pubCon.publish("draintest", null);
            pubCon.flush(Duration.ofSeconds(1));
            subCon.flush(Duration.ofSeconds(1));

            // Drain will unsub the dispatcher, only messages that already arrived
            // are there
            CompletableFuture<Boolean> tracker = d.drain(Duration.ofSeconds(4));

            assertTrue(tracker.get(5, TimeUnit.SECONDS)); // wait for the drain to complete
            assertEquals(count.get(), 4); // Should get both, two times.
            assertFalse(d.isActive());
            assertEquals(((NatsConnection) subCon).getConsumerCount(), 0);
        }
    }

    @Test
    public void testSimpleConnectionDrain() throws Exception {
        try (NatsTestServer ts = new NatsTestServer(false);
                Connection subCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build());
                Connection pubCon = Nats.connect(new Options.Builder().server(ts.getURI()).maxReconnects(0).build())) {
            assertTrue("Connected Status", Connection.Status.CONNECTED == subCon.getStatus());
            assertTrue("Connected Status", Connection.Status.CONNECTED == pubCon.getStatus());

            AtomicInteger count = new AtomicInteger();
            Dispatcher d = subCon.createDispatcher((msg) -> {
                count.incrementAndGet();
                try {
{'fraction_non_alphanumeric': 0.11693511693511693, 'fraction_numerical': 0.007098007098007098, 'mean_word_length': 2.891643059490085, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7653397', 'n_tokens_mistral': 2960, 'n_tokens_neox': 2872, 'n_words': 700}