<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.pulsar.tests.integration.functions;

import lombok.Cleanup;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.pulsar.client.admin.PulsarAdmin;
import org.apache.pulsar.client.admin.PulsarAdminException;
import org.apache.pulsar.client.api.Consumer;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.Producer;
import org.apache.pulsar.client.api.PulsarClient;
import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.api.SubscriptionType;
import org.apache.pulsar.common.functions.FunctionState;
import org.apache.pulsar.common.policies.data.SinkStatus;
import org.apache.pulsar.common.policies.data.SourceStatus;
import org.apache.pulsar.tests.integration.docker.ContainerExecException;
import org.apache.pulsar.tests.integration.docker.ContainerExecResult;
import org.apache.pulsar.tests.integration.functions.utils.CommandGenerator;
import org.apache.pulsar.tests.integration.functions.utils.CommandGenerator.Runtime;
import org.apache.pulsar.tests.integration.suites.PulsarStandaloneTestSuite;
import org.apache.pulsar.tests.integration.topologies.PulsarCluster;
import org.testng.annotations.Test;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.apache.pulsar.tests.integration.functions.utils.CommandGenerator.JAVAJAR;
import static org.apache.pulsar.tests.integration.suites.PulsarTestSuite.retryStrategically;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;

/**
 * State related test cases.
 */
@Slf4j
public class PulsarStateTest extends PulsarStandaloneTestSuite {

    public static final String WORDCOUNT_PYTHON_CLASS =
        "wordcount_function.WordCountFunction";

    public static final String WORDCOUNT_PYTHON_FILE = "wordcount_function.py";

    @Test(groups = {"python_state", "state", "function", "python_function"})
    public void testPythonWordCountFunction() throws Exception {
        String inputTopicName = "test-wordcount-py-input-" + randomName(8);
        String outputTopicName = "test-wordcount-py-output-" + randomName(8);
        String functionName = "test-wordcount-py-fn-" + randomName(8);

        final int numMessages = 10;

        // submit the exclamation function
        submitExclamationFunction(
            Runtime.PYTHON, inputTopicName, outputTopicName, functionName);

        // get function info
        getFunctionInfoSuccess(functionName);

        // publish and consume result
        publishAndConsumeMessages(inputTopicName, outputTopicName, numMessages);

        // get function status
        getFunctionStatus(functionName, numMessages);

        // get state
        queryState(functionName, "hello", numMessages);
        queryState(functionName, "test", numMessages);
        for (int i = 0; i < numMessages; i++) {
            queryState(functionName, "message-" + i, 1);
        }

        // delete function
        deleteFunction(functionName);

        // get function info
        getFunctionInfoNotFound(functionName);
    }

    @Test(groups = {"java_state", "state", "function", "java_function"})
    public void testSourceState() throws Exception {
        String outputTopicName = "test-state-source-output-" + randomName(8);
        String sourceName = "test-state-source-" + randomName(8);

        submitSourceConnector(sourceName, outputTopicName, "org.apache.pulsar.tests.integration.io.TestStateSource",  JAVAJAR);

        // get source info
        getSourceInfoSuccess(sourceName);

        // get source status
        getSourceStatus(sourceName);

        try (PulsarAdmin admin = PulsarAdmin.builder().serviceHttpUrl(container.getHttpServiceUrl()).build()) {

            retryStrategically((test) -> {
                try {
                    SourceStatus status = admin.sources().getSourceStatus("public", "default", sourceName);
                    return status.getInstances().size() > 0 && status.getInstances().get(0).getStatus().numWritten > 0;
                } catch (PulsarAdminException e) {
                    return false;
                }
            }, 10, 200);

            SourceStatus status = admin.sources().getSourceStatus("public", "default", sourceName);
            assertEquals(status.getInstances().size(), 1);
            assertTrue(status.getInstances().get(0).getStatus().numWritten > 0);

            FunctionState functionState = admin.functions().getFunctionState("public", "default", sourceName, "initial");
            assertEquals(functionState.getStringValue(), "val1");

            functionState = admin.functions().getFunctionState("public", "default", sourceName, "now");
            assertTrue(functionState.getStringValue().matches("val1-.*"));
        }

        // delete source
        deleteSource(sourceName);

        getSourceInfoNotFound(sourceName);
    }

    @Test(groups = {"java_state", "state", "function", "java_function"})
    public void testSinkState() throws Exception {
        String inputTopicName = "test-state-sink-input-" + randomName(8);
        String sinkName = "test-state-sink-" + randomName(8);
        int numMessages = 10;

        submitSinkConnector(sinkName, inputTopicName, "org.apache.pulsar.tests.integration.io.TestStateSink",  JAVAJAR);

        // get sink info
        getSinkInfoSuccess(sinkName);

        // get sink status
        getSinkStatus(sinkName);

        try (PulsarAdmin admin = PulsarAdmin.builder().serviceHttpUrl(container.getHttpServiceUrl()).build()) {

            // java supports schema
            @Cleanup PulsarClient client = PulsarClient.builder()
                    .serviceUrl(container.getPlainTextServiceUrl())
                    .build();
            @Cleanup Producer<String> producer = client.newProducer(Schema.STRING)
                    .topic(inputTopicName)
                    .create();

            FunctionState functionState = admin.functions().getFunctionState("public", "default", sinkName, "initial");
            assertEquals(functionState.getStringValue(), "val1");

            for (int i = 0; i < numMessages; i++) {
                producer.send("foo");
            }

            retryStrategically((test) -> {
                try {
                    SinkStatus status = admin.sinks().getSinkStatus("public", "default", sinkName);
                    return status.getInstances().size() > 0 && status.getInstances().get(0).getStatus().numWrittenToSink > 0;
                } catch (PulsarAdminException e) {
                    return false;
                }
            }, 10, 200);

            SinkStatus status = admin.sinks().getSinkStatus("public", "default", sinkName);
            assertEquals(status.getInstances().size(), 1);
            assertTrue(status.getInstances().get(0).getStatus().numWrittenToSink > 0);

            functionState = admin.functions().getFunctionState("public", "default", sinkName, "now");
            assertEquals(functionState.getStringValue(), String.format("val1-%d", numMessages - 1));<commit_msg>[Flaky test] Attempt to fix flakiness in PulsarStateTest (#9870)<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.pulsar.tests.integration.functions;

import lombok.Cleanup;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.pulsar.client.admin.PulsarAdmin;
import org.apache.pulsar.client.admin.PulsarAdminException;
import org.apache.pulsar.client.api.Consumer;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.Producer;
import org.apache.pulsar.client.api.PulsarClient;
import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.api.SubscriptionType;
import org.apache.pulsar.common.functions.FunctionState;
import org.apache.pulsar.common.policies.data.SinkStatus;
import org.apache.pulsar.common.policies.data.SourceStatus;
import org.apache.pulsar.tests.integration.docker.ContainerExecException;
import org.apache.pulsar.tests.integration.docker.ContainerExecResult;
import org.apache.pulsar.tests.integration.functions.utils.CommandGenerator;
import org.apache.pulsar.tests.integration.functions.utils.CommandGenerator.Runtime;
import org.apache.pulsar.tests.integration.suites.PulsarStandaloneTestSuite;
import org.apache.pulsar.tests.integration.topologies.PulsarCluster;
import org.awaitility.Awaitility;
import org.testng.annotations.Test;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.apache.pulsar.tests.integration.functions.utils.CommandGenerator.JAVAJAR;
import static org.apache.pulsar.tests.integration.suites.PulsarTestSuite.retryStrategically;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;

/**
 * State related test cases.
 */
@Slf4j
public class PulsarStateTest extends PulsarStandaloneTestSuite {

    public static final String WORDCOUNT_PYTHON_CLASS =
        "wordcount_function.WordCountFunction";

    public static final String WORDCOUNT_PYTHON_FILE = "wordcount_function.py";

    @Test(groups = {"python_state", "state", "function", "python_function"})
    public void testPythonWordCountFunction() throws Exception {
        String inputTopicName = "test-wordcount-py-input-" + randomName(8);
        String outputTopicName = "test-wordcount-py-output-" + randomName(8);
        String functionName = "test-wordcount-py-fn-" + randomName(8);

        final int numMessages = 10;

        // submit the exclamation function
        submitExclamationFunction(
            Runtime.PYTHON, inputTopicName, outputTopicName, functionName);

        // get function info
        getFunctionInfoSuccess(functionName);

        // publish and consume result
        publishAndConsumeMessages(inputTopicName, outputTopicName, numMessages);

        // get function status
        getFunctionStatus(functionName, numMessages);

        // get state
        queryState(functionName, "hello", numMessages);
        queryState(functionName, "test", numMessages);
        for (int i = 0; i < numMessages; i++) {
            queryState(functionName, "message-" + i, 1);
        }

        // delete function
        deleteFunction(functionName);

        // get function info
        getFunctionInfoNotFound(functionName);
    }

    @Test(groups = {"java_state", "state", "function", "java_function"})
    public void testSourceState() throws Exception {
        String outputTopicName = "test-state-source-output-" + randomName(8);
        String sourceName = "test-state-source-" + randomName(8);

        submitSourceConnector(sourceName, outputTopicName, "org.apache.pulsar.tests.integration.io.TestStateSource",  JAVAJAR);

        // get source info
        getSourceInfoSuccess(sourceName);

        // get source status
        getSourceStatus(sourceName);

        try (PulsarAdmin admin = PulsarAdmin.builder().serviceHttpUrl(container.getHttpServiceUrl()).build()) {

            Awaitility.await().ignoreExceptions().untilAsserted(() -> {
                SourceStatus status = admin.sources().getSourceStatus("public", "default", sourceName);
                assertEquals(status.getInstances().size(), 1);
                assertTrue(status.getInstances().get(0).getStatus().numWritten > 0);
            });

            {
                FunctionState functionState =
                        admin.functions().getFunctionState("public", "default", sourceName, "initial");
                assertEquals(functionState.getStringValue(), "val1");
            }

            Awaitility.await().ignoreExceptions().untilAsserted(() -> {
                FunctionState functionState = admin.functions().getFunctionState("public", "default", sourceName, "now");
                assertTrue(functionState.getStringValue().matches("val1-.*"));
            });
        }

        // delete source
        deleteSource(sourceName);

        getSourceInfoNotFound(sourceName);
    }

    @Test(groups = {"java_state", "state", "function", "java_function"})
    public void testSinkState() throws Exception {
        String inputTopicName = "test-state-sink-input-" + randomName(8);
        String sinkName = "test-state-sink-" + randomName(8);
        int numMessages = 10;

        submitSinkConnector(sinkName, inputTopicName, "org.apache.pulsar.tests.integration.io.TestStateSink",  JAVAJAR);

        // get sink info
        getSinkInfoSuccess(sinkName);

        // get sink status
        getSinkStatus(sinkName);

        try (PulsarAdmin admin = PulsarAdmin.builder().serviceHttpUrl(container.getHttpServiceUrl()).build()) {

            // java supports schema
            @Cleanup PulsarClient client = PulsarClient.builder()
                    .serviceUrl(container.getPlainTextServiceUrl())
                    .build();
            @Cleanup Producer<String> producer = client.newProducer(Schema.STRING)
                    .topic(inputTopicName)
                    .create();

            {
                FunctionState functionState =
                        admin.functions().getFunctionState("public", "default", sinkName, "initial");
                assertEquals(functionState.getStringValue(), "val1");
            }

            for (int i = 0; i < numMessages; i++) {
                producer.send("foo");
            }

            Awaitility.await().ignoreExceptions().untilAsserted(() -> {
                SinkStatus status = admin.sinks().getSinkStatus("public", "default", sinkName);
                assertEquals(status.getInstances().size(), 1);
                assertTrue(status.getInstances().get(0).getStatus().numWrittenToSink > 0);
            });

            Awaitility.await().ignoreExceptions().untilAsserted(() -> {
                FunctionState functionState = admin.functions().getFunctionState("public", "default", sinkName, "now");
                assertEquals(functionState.getStringValue(), String.format("val1-%d", numMessages - 1));
            });
{'fraction_non_alphanumeric': 0.11158743361699405, 'fraction_numerical': 0.005630558577004287, 'mean_word_length': 3.1185770750988144, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 19, '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': '3146368', 'n_tokens_mistral': 4320, 'n_tokens_neox': 4152, 'n_words': 895}