/*
 * Copyright 2021 PingCAP, Inc.
 *
 * 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,
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"context"
	"os"
	"strings"
	"sync"
	"testing"

	"github.com/segmentio/kafka-go"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
)

type MockMapper struct {
	pipe chan map[string]interface{}
}

func NewMockMapper() *MockMapper {
	return &MockMapper{pipe: make(chan map[string]interface{})}
}

func (m *MockMapper) Handle(ctx context.Context, data map[string]interface{}) error {
	m.pipe <- data
	return nil
}

func TestKafkaConsumer_Start(t *testing.T) {
	broker := strings.Split(os.Getenv("BROKER"), ",")
	topic := os.Getenv("TOPIC")
	groupID := os.Getenv("GROUP_ID")
	c, err := NewKafkaConsumer(kafka.DefaultDialer, broker, topic, groupID)
	if err != nil {
		t.Fatal(err)
	}
	m := NewMockMapper()
	c.SetMapper(m)
	ctx := context.Background()
	wg := sync.WaitGroup{}
	wg.Add(2)
	go func() {
		defer wg.Done()
		if err := c.Start(ctx); err != nil {
			t.Error(err)
		}
		close(m.pipe)
	}()
	go func() {
		defer wg.Done()
		for {
			select {
			case data, ok := <-m.pipe:
				if !ok {
					return
				}
				t.Log(data)
			}
		}
	}()
	wg.Wait()
}

func TestSQSConsumer_Start(t *testing.T) {
	sqsURL := os.Getenv("SQS_URL")
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))
	c, err := NewSQSConsumer(sqs.New(sess), sqsURL,
		WithVisibilityTimeout(20),
	)
	if err != nil {
		t.Fatal(err)
	}
	m := NewMockMapper()
	c.SetMapper(m)
	ctx := context.Background()
	wg := sync.WaitGroup{}
	wg.Add(2)
	go func() {
		defer wg.Done()
		if err := c.Start(ctx); err != nil {
			t.Error(err)
		}
		close(m.pipe)
	}()
	go func() {
		defer wg.Done()
		for {
			select {
			case data, ok := <-m.pipe:
				if !ok {
					return
				}
				t.Log(data)
			}
		}
	}()
	wg.Wait()
}