package audit

import (
	"context"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/stretchr/testify/assert"
)

func TestDeleteBatchStream(t *testing.T) {
	t.Run("less than one batch", func(t *testing.T) {
		batches := deleteBatchStream{
			Size:     10,
			IDStream: &arrayIDStream{[]string{"a"}},
		}

		var (
			b    []*dynamodb.WriteRequest
			last bool
		)

		b, last = batches.Next(context.Background())
		assert.Equal(t, []*dynamodb.WriteRequest{
			{
				DeleteRequest: &dynamodb.DeleteRequest{
					Key: map[string]*dynamodb.AttributeValue{
						"uuid": {S: aws.String("a")},
					},
				},
			},
		}, b)
		assert.True(t, last)

		b, last = batches.Next(context.Background())
		assert.Empty(t, b)
		assert.True(t, last)
	})

	t.Run("exactly one batch", func(t *testing.T) {
		batches := deleteBatchStream{
			Size:     2,
			IDStream: &arrayIDStream{[]string{"a", "b"}},
		}

		var (
			b    []*dynamodb.WriteRequest
			last bool
		)

		b, last = batches.Next(context.Background())
		assert.Equal(t, []*dynamodb.WriteRequest{
			{
				DeleteRequest: &dynamodb.DeleteRequest{
					Key: map[string]*dynamodb.AttributeValue{
						"uuid": {S: aws.String("a")},
					},
				},
			},
			{
				DeleteRequest: &dynamodb.DeleteRequest{
					Key: map[string]*dynamodb.AttributeValue{
						"uuid": {S: aws.String("b")},
					},
				},
			},
		}, b)
		assert.True(t, last)

		b, last = batches.Next(context.Background())
		assert.Empty(t, b)
		assert.True(t, last)
	})
}
