Kiro/internal/storage/dynamodb_test.go

package storage

import (
	"context"
	"os"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
)

func setupTestClient(t *testing.T) (*DynamoDBClient, context.Context) {
	t.Helper()

	// Set dummy AWS credentials for local DynamoDB
	os.Setenv("AWS_ACCESS_KEY_ID", "dummy")
	os.Setenv("AWS_SECRET_ACCESS_KEY", "dummy")
	os.Setenv("AWS_REGION", "us-east-1")

	endpoint := os.Getenv("DYNAMODB_ENDPOINT")
	if endpoint == "" {
		endpoint = "http://localhost:8000"
	}

	ctx := context.Background()
	client, err := NewDynamoDBClient(ctx, endpoint)
	if err != nil {
		t.Skipf("Skipping test: DynamoDB not available: %v", err)
	}

	// Test connection by listing tables
	_, err = client.client.ListTables(ctx, &dynamodb.ListTablesInput{})
	if err != nil {
		t.Skipf("Skipping test: Cannot connect to DynamoDB: %v", err)
	}

	return client, ctx
}

func TestNewDynamoDBClient(t *testing.T) {
	client, _ := setupTestClient(t)

	if client == nil {
		t.Fatal("Expected non-nil client")
	}

	if client.client == nil {
		t.Fatal("Expected non-nil underlying client")
	}
}

func TestTransactWriteItems(t *testing.T) {
	client, ctx := setupTestClient(t)

	// Create test table
	tableName := "TestTransactions"
	createTestTable(t, ctx, client, tableName)
	defer deleteTestTable(t, ctx, client, tableName)

	// Test transaction with two puts
	testID1 := "test-txn-1"
	testID2 := "test-txn-2"

	input := &dynamodb.TransactWriteItemsInput{
		TransactItems: []types.TransactWriteItem{
			{
				Put: &types.Put{
					TableName: aws.String(tableName),
					Item: map[string]types.AttributeValue{
						"id":    &types.AttributeValueMemberS{Value: testID1},
						"value": &types.AttributeValueMemberS{Value: "value1"},
					},
				},
			},
			{
				Put: &types.Put{
					TableName: aws.String(tableName),
					Item: map[string]types.AttributeValue{
						"id":    &types.AttributeValueMemberS{Value: testID2},
						"value": &types.AttributeValueMemberS{Value: "value2"},
					},
				},
			},
		},
	}

	err := client.TransactWriteItems(ctx, input)
	if err != nil {
		t.Fatalf("TransactWriteItems failed: %v", err)
	}

	// Verify both items were written
	output1, err := client.GetItem(ctx, &dynamodb.GetItemInput{
		TableName: aws.String(tableName),
		Key: map[string]types.AttributeValue{
			"id": &types.AttributeValueMemberS{Value: testID1},
		},
	})
	if err != nil {
		t.Fatalf("GetItem failed: %v", err)
	}
	if len(output1.Item) == 0 {
		t.Fatal("Expected item to exist after transaction")
	}

	output2, err := client.GetItem(ctx, &dynamodb.GetItemInput{
		TableName: aws.String(tableName),
		Key: map[string]types.AttributeValue{
			"id": &types.AttributeValueMemberS{Value: testID2},
		},
	})
	if err != nil {
		t.Fatalf("GetItem failed: %v", err)
	}
	if len(output2.Item) == 0 {
		t.Fatal("Expected item to exist after transaction")
	}
}

func TestBatchGetItems(t *testing.T) {
	client, ctx := setupTestClient(t)

	// Create test table
	tableName := "TestBatchGet"
	createTestTable(t, ctx, client, tableName)
	defer deleteTestTable(t, ctx, client, tableName)

	// Put test items
	testIDs := []string{"batch-1", "batch-2", "batch-3"}
	for _, id := range testIDs {
		err := client.PutItem(ctx, &dynamodb.PutItemInput{
			TableName: aws.String(tableName),
			Item: map[string]types.AttributeValue{
				"id":    &types.AttributeValueMemberS{Value: id},
				"value": &types.AttributeValueMemberS{Value: "test-value"},
			},
		})
		if err != nil {
			t.Fatalf("PutItem failed: %v", err)
		}
	}

	// Batch get items
	keys := make([]map[string]types.AttributeValue, len(testIDs))
	for i, id := range testIDs {
		keys[i] = map[string]types.AttributeValue{
			"id": &types.AttributeValueMemberS{Value: id},
		}
	}

	output, err := client.BatchGetItems(ctx, &dynamodb.BatchGetItemInput{
		RequestItems: map[string]types.KeysAndAttributes{
			tableName: {
				Keys: keys,
			},
		},
	})
	if err != nil {
		t.Fatalf("BatchGetItems failed: %v", err)
	}

	if len(output.Responses[tableName]) != len(testIDs) {
		t.Fatalf("Expected %d items, got %d", len(testIDs), len(output.Responses[tableName]))
	}
}

func TestBatchWriteItems(t *testing.T) {
	client, ctx := setupTestClient(t)

	// Create test table
	tableName := "TestBatchWrite"
	createTestTable(t, ctx, client, tableName)
	defer deleteTestTable(t, ctx, client, tableName)

	// Batch write items
	testIDs := []string{"write-1", "write-2", "write-3"}
	writeRequests := make([]types.WriteRequest, len(testIDs))
	for i, id := range testIDs {
		writeRequests[i] = types.WriteRequest{
			PutRequest: &types.PutRequest{
				Item: map[string]types.AttributeValue{
					"id":    &types.AttributeValueMemberS{Value: id},
					"value": &types.AttributeValueMemberS{Value: "batch-value"},
				},
			},
		}
	}

	err := client.BatchWriteItems(ctx, &dynamodb.BatchWriteItemInput{
		RequestItems: map[string][]types.WriteRequest{
			tableName: writeRequests,
		},
	})
	if err != nil {
		t.Fatalf("BatchWriteItems failed: %v", err)
	}

	// Verify items were written
	for _, id := range testIDs {
		output, err := client.GetItem(ctx, &dynamodb.GetItemInput{
			TableName: aws.String(tableName),
			Key: map[string]types.AttributeValue{
				"id": &types.AttributeValueMemberS{Value: id},
			},
		})
		if err != nil {
			t.Fatalf("GetItem failed: %v", err)
		}
		if len(output.Item) == 0 {
			t.Fatalf("Expected item %s to exist after batch write", id)
		}
	}
}

func TestPutAndGetItem(t *testing.T) {
	client, ctx := setupTestClient(t)

	// Create test table
	tableName := "TestPutGet"
	createTestTable(t, ctx, client, tableName)
	defer deleteTestTable(t, ctx, client, tableName)

	// Put item
	testID := "put-get-test"
	err := client.PutItem(ctx, &dynamodb.PutItemInput{
		TableName: aws.String(tableName),
		Item: map[string]types.AttributeValue{
			"id":    &types.AttributeValueMemberS{Value: testID},
			"value": &types.AttributeValueMemberS{Value: "test-value"},
		},
	})
	if err != nil {
		t.Fatalf("PutItem failed: %v", err)
	}

	// Get item
	output, err := client.GetItem(ctx, &dynamodb.GetItemInput{
		TableName: aws.String(tableName),
		Key: map[string]types.AttributeValue{
			"id": &types.AttributeValueMemberS{Value: testID},
		},
	})
	if err != nil {
		t.Fatalf("GetItem failed: %v", err)
	}

	if len(output.Item) == 0 {
		t.Fatal("Expected item to exist")
	}

	valueAttr, ok := output.Item["value"]
	if !ok {
		t.Fatal("Expected 'value' attribute")
	}

	value := valueAttr.(*types.AttributeValueMemberS).Value
	if value != "test-value" {
		t.Fatalf("Expected value 'test-value', got '%s'", value)
	}
}

func TestUpdateItem(t *testing.T) {
	client, ctx := setupTestClient(t)

	// Create test table
	tableName := "TestUpdate"
	createTestTable(t, ctx, client, tableName)
	defer deleteTestTable(t, ctx, client, tableName)

	// Put initial item
	testID := "update-test"
	err := client.PutItem(ctx, &dynamodb.PutItemInput{
		TableName: aws.String(tableName),
		Item: map[string]types.AttributeValue{
			"id":    &types.AttributeValueMemberS{Value: testID},
			"value": &types.AttributeValueMemberS{Value: "initial"},
		},
	})
	if err != nil {
		t.Fatalf("PutItem failed: %v", err)
	}

	// Update item
	_, err = client.UpdateItem(ctx, &dynamodb.UpdateItemInput{
		TableName: aws.String(tableName),
		Key: map[string]types.AttributeValue{
			"id": &types.AttributeValueMemberS{Value: testID},
		},
		UpdateExpression: aws.String("SET #v = :val"),
		ExpressionAttributeNames: map[string]string{
			"#v": "value",
		},
		ExpressionAttributeValues: map[string]types.AttributeValue{
			":val": &types.AttributeValueMemberS{Value: "updated"},
		},
	})
	if err != nil {
		t.Fatalf("UpdateItem failed: %v", err)
	}

	// Verify update
	output, err := client.GetItem(ctx, &dynamodb.GetItemInput{
		TableName: aws.String(tableName),
		Key: map[string]types.AttributeValue{
			"id": &types.AttributeValueMemberS{Value: testID},
		},
	})
	if err != nil {
		t.Fatalf("GetItem failed: %v", err)
	}

	value := output.Item["value"].(*types.AttributeValueMemberS).Value
	if value != "updated" {
		t.Fatalf("Expected value 'updated', got '%s'", value)
	}
}

func TestDeleteItem(t *testing.T) {
	client, ctx := setupTestClient(t)

	// Create test table
	tableName := "TestDelete"
	createTestTable(t, ctx, client, tableName)
	defer deleteTestTable(t, ctx, client, tableName)

	// Put item
	testID := "delete-test"
	err := client.PutItem(ctx, &dynamodb.PutItemInput{
		TableName: aws.String(tableName),
		Item: map[string]types.AttributeValue{
			"id":    &types.AttributeValueMemberS{Value: testID},
			"value": &types.AttributeValueMemberS{Value: "test"},
		},
	})
	if err != nil {
		t.Fatalf("PutItem failed: %v", err)
	}

	// Delete item
	err = client.DeleteItem(ctx, &dynamodb.DeleteItemInput{
		TableName: aws.String(tableName),
		Key: map[string]types.AttributeValue{
			"id": &types.AttributeValueMemberS{Value: testID},
		},
	})
	if err != nil {
		t.Fatalf("DeleteItem failed: %v", err)
	}

	// Verify deletion
	output, err := client.GetItem(ctx, &dynamodb.GetItemInput{
		TableName: aws.String(tableName),
		Key: map[string]types.AttributeValue{
			"id": &types.AttributeValueMemberS{Value: testID},
		},
	})
	if err != nil {
		t.Fatalf("GetItem failed: %v", err)
	}

	if len(output.Item) != 0 {
		t.Fatal("Expected item to be deleted")
	}
}

func TestQuery(t *testing.T) {
	client, ctx := setupTestClient(t)

	// Create test table
	tableName := "TestQuery"
	createTestTable(t, ctx, client, tableName)
	defer deleteTestTable(t, ctx, client, tableName)

	// Put test items
	testIDs := []string{"query-1", "query-2", "query-3"}
	for _, id := range testIDs {
		err := client.PutItem(ctx, &dynamodb.PutItemInput{
			TableName: aws.String(tableName),
			Item: map[string]types.AttributeValue{
				"id":    &types.AttributeValueMemberS{Value: id},
				"value": &types.AttributeValueMemberS{Value: "test"},
			},
		})
		if err != nil {
			t.Fatalf("PutItem failed: %v", err)
		}
	}

	// Query for specific item
	output, err := client.Query(ctx, &dynamodb.QueryInput{
		TableName:              aws.String(tableName),
		KeyConditionExpression: aws.String("id = :id"),
		ExpressionAttributeValues: map[string]types.AttributeValue{
			":id": &types.AttributeValueMemberS{Value: "query-1"},
		},
	})
	if err != nil {
		t.Fatalf("Query failed: %v", err)
	}

	if len(output.Items) != 1 {
		t.Fatalf("Expected 1 item, got %d", len(output.Items))
	}
}

// Helper functions

func createTestTable(t *testing.T, ctx context.Context, client *DynamoDBClient, tableName string) {
	t.Helper()

	_, err := client.client.CreateTable(ctx, &dynamodb.CreateTableInput{
		TableName: aws.String(tableName),
		AttributeDefinitions: []types.AttributeDefinition{
			{
				AttributeName: aws.String("id"),
				AttributeType: types.ScalarAttributeTypeS,
			},
		},
		KeySchema: []types.KeySchemaElement{
			{
				AttributeName: aws.String("id"),
				KeyType:       types.KeyTypeHash,
			},
		},
		BillingMode: types.BillingModePayPerRequest,
	})
	if err != nil {
		t.Fatalf("Failed to create test table: %v", err)
	}

	// Wait for table to be active
	waiter := dynamodb.NewTableExistsWaiter(client.client)
	err = waiter.Wait(ctx, &dynamodb.DescribeTableInput{
		TableName: aws.String(tableName),
	}, 30*time.Second)
	if err != nil {
		t.Fatalf("Failed waiting for table to be active: %v", err)
	}
}

func deleteTestTable(t *testing.T, ctx context.Context, client *DynamoDBClient, tableName string) {
	t.Helper()

	_, err := client.client.DeleteTable(ctx, &dynamodb.DeleteTableInput{
		TableName: aws.String(tableName),
	})
	if err != nil {
		t.Logf("Warning: Failed to delete test table: %v", err)
	}
}