java/feathub-connectors/flink-connectors/flink-connector-redis/src/test/java/org/apache/flink/streaming/connectors/redis/RedisSinkITCase.java
/*
 * Copyright 2022 The FeatHub 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
 *
 *     https://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.flink.streaming.connectors.redis;

import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.api.Schema;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.TableDescriptor;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.types.Row;

import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import redis.clients.jedis.Jedis;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import static org.assertj.core.api.Assertions.assertThat;

/** Tests redis sink. */
public class RedisSinkITCase extends RedisITCaseBase {

    private static final int NUM_ELEMENTS = 20;

    private StreamExecutionEnvironment env;

    private StreamTableEnvironment tEnv;

    private Jedis jedis;

    @Before
    public void setUp() {
        jedis = new Jedis(REDIS_HOST, redisPort.getPort());
        env = StreamExecutionEnvironment.getExecutionEnvironment();
        tEnv = StreamTableEnvironment.create(env);
    }

    private void buildAndExecute(DataStream<Row> stream, Map<String, Object> configs)
            throws Exception {
        Table table = tEnv.fromDataStream(stream);

        TableDescriptor.Builder builder =
                TableDescriptor.forConnector("redis")
                        .schema(
                                Schema.newBuilder()
                                        .column("__KEY__f1", DataTypes.STRING())
                                        .column("f1", DataTypes.STRING())
                                        .build());

        for (Map.Entry<String, Object> entry : configs.entrySet()) {
            builder = builder.option(entry.getKey(), String.valueOf(entry.getValue()));
        }

        tEnv.createTemporaryTable("redis_sink", builder.build());

        table.executeInsert("redis_sink").await(30, TimeUnit.SECONDS);
    }

    private void verifyOutputResult() {
        assertThat(jedis.keys("*")).hasSize(NUM_ELEMENTS);

        for (int i = 1; i <= NUM_ELEMENTS; i++) {
            assertThat(jedis.get("test_namespace:" + i + ":f1")).isEqualTo(Integer.toString(i));
        }
    }

    @Test
    public void testWriteToRedis() throws Exception {
        DataStream<Row> stream =
                env.fromSequence(1, NUM_ELEMENTS)
                        .map(
                                x -> Row.of("test_namespace:" + x + ":f1", x.toString()),
                                new RowTypeInfo(Types.STRING, Types.STRING));

        Map<String, Object> configs = new HashMap<>();
        configs.put("host", REDIS_HOST);
        configs.put("port", redisPort.getPort());
        configs.put("dbNum", 0);

        buildAndExecute(stream, configs);
        verifyOutputResult();
    }

    @Test
    public void testChangeDBNum() throws Exception {
        int dbNum = 2;

        DataStream<Row> stream =
                env.fromSequence(1, NUM_ELEMENTS)
                        .map(
                                x -> Row.of("test_namespace:" + x + ":f1", x.toString()),
                                new RowTypeInfo(Types.STRING, Types.STRING));
        Map<String, Object> configs = new HashMap<>();
        configs.put("host", REDIS_HOST);
        configs.put("port", redisPort.getPort());
        configs.put("dbNum", dbNum);

        buildAndExecute(stream, configs);

        assertThat(jedis.keys("*")).isEmpty();
        jedis.select(dbNum);
        verifyOutputResult();
    }

    @After
    public void tearDown() {
        for (String key : jedis.keys("*")) {
            jedis.del(key);
        }

        if (jedis != null) {
            jedis.close();
        }
    }
}
