Go - aws-sdk-go/service Unit Test

package repository

import (
    "context"
    "fmt"
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/dynamodb"
    "github.com/aws/aws-sdk-go/service/dynamodb/dynamodbattribute"
    // ...
)

// SDK wrapper interface, which can be able to unit test.
type DynamoDbWrapper interface {
    // Define the same as SDK.
    Query(input *dynamodb.QueryInput) (*dynamodb.QueryOutput, error)

    // Define, if it needs more.
    // ...     
}

// Default implementation.
// Use this if the argument of "NewDynamoDbRepository" is not passed.
type DynamoDbWrapperAdapter struct {
    svc *dynamodb.DynamoDB
}

// Dispatch simply.
func (d *DynamoDbWrapperAdapter) Query(input *dynamodb.QueryInput) (*dynamodb.QueryOutput, error) {
    return d.svc.Query(input)
}

// Set wrp to null. In case unit test, set mock interface.
func NewDynamoDbRepository(wrp DynamoDbWrapper) updatetimerevent.Repository {
    if wrp == nil {
        wrp = &DynamoDbWrapperAdapter{
            svc: dynamodb.New(session.New()),
        }
    }
    return &DynamoDbRepository{
        wrp: wrp,
    }
}

// Find something.
func (r *DynamoDbRepository) FindSome(ctx context.Context, userId string) (some *SomeStruct, err error) {
    input := &dynamodb.QueryInput{
        ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
            ":userid": {
                S: aws.String(userId),
            },
        },
        KeyConditionExpression: aws.String("UserId = :userid"),
        TableName:              aws.String("DYNAMODB_TABLE"),
    }

    // Call mock method, if r.wrp is set mock.
    result, err := r.wrp.Query(input)
    if err != nil {
        return
    }
    // ...
}

t.Run("ok:Query", func(t *testing.T) {
    caseUserId := "test user"
    caseInput := &dynamodb.QueryInput{
        ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
            ":userid": {
                S: aws.String(caseUserId),
            },
        },
        KeyConditionExpression: aws.String("UserId = :userid"),
        TableName:              aws.String("DYNAMODB_TABLE"),
    }
    caseItem := &dynamodb.QueryOutput{
        Items: []map[string]*dynamodb.AttributeValue{},
    }

    ctrl := gomock.NewController(t)
    defer ctrl.Finish()

    s := NewMockDynamoDbWrapper(ctrl)
    s.EXPECT().Query(gomock.Eq(caseInput)).Return(caseItem, nil)

    repo := NewDynamoDbRepository(s)
    got, err := repo.FindSome(context.TODO(), caseUserId)
    assert.NoError(t, err)
    // assert...
})