aws-doc-sdk-examples/gov2/bedrock-runtime/stubs/invoke_model_stubs.go at bc5d30afcff873a00fd6c057a1faa0c87f29de5f · rlhagerm/aws-doc-sdk-examples · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

// Defines stubs used for unit testing the Bedrock Runtime actions.

package stubs

import (
	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
	"github.com/awsdocs/aws-doc-sdk-examples/gov2/testtools"
)

type StubInvokeModelParams struct {
    Request  []byte
    Response []byte
    ModelId  string
    RaiseErr *testtools.StubError
}

func StubInvokeModel(params StubInvokeModelParams) testtools.Stub {
    return testtools.Stub{
    		OperationName: "InvokeModel",
    		Input:	&bedrockruntime.InvokeModelInput{
    			Body:        params.Request,
    			ModelId:     aws.String(params.ModelId),
    			ContentType: aws.String("application/json"),
    		},
    		Output:	&bedrockruntime.InvokeModelOutput{ Body: params.Response, },
    		Error:	params.RaiseErr,
    	}
}

// func StubInvokeTitanImage(requestBytes []byte, raiseErr *testtools.StubError) testtools.Stub {
//     fakeTitanImageResponse := actions.TitanImageResponse{
//         Images: []string {
//             "FakeBase64String==",
//         },
//     }
//
//     responseBytes, err := json.Marshal(fakeTitanImageResponse)
//     if err != nil {
//         panic(err)
//     }
//
//     return testtools.Stub{
//         OperationName: "InvokeModel",
//         Input:	&bedrockruntime.InvokeModelInput{
//             Body:        requestBytes,
//             ModelId:     aws.String("amazon.titan-image-generator-v2:0"),
//             ContentType: aws.String("application/json"),
//         },
//         Output:	&bedrockruntime.InvokeModelOutput{
//             Body:        responseBytes,
//         },
//         Error:	raiseErr,
//     }
// }