Skip to content

Commit f95900f

Browse files
authored
feat(azure): add cosmosdb module (#3452)
* feat: add cosmosdb module * fix: typo * refactor: move cosmosdb package to azure module * chore: remove references to cosmosdb module * docs: update docs/modules/azure.md * feat: improve readiness check with port listening * docs: improve ContainerPolicy description * fix: wrap error in NewContainerPolicy * docs: document hardcoded key * chore(deps): bump azcore version to "v1.19.1" * fix: unnecessary nested wait.ForAll * docs: update test account key comment * fix: .vscode/.testcontainers-go.code-workspace extra line * chore: doc ConnectionString and move const to top
1 parent 1ed2735 commit f95900f

File tree

7 files changed

+341
-15
lines changed

7 files changed

+341
-15
lines changed

docs/modules/azure.md

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ The Azure module exposes the following Go packages:
2323
- [ServiceBus](#servicebus): `github.com/testcontainers/testcontainers-go/modules/azure/servicebus`.
2424
!!! warning "EULA Acceptance"
2525
Due to licensing restrictions you are required to explicitly accept an End User License Agreement (EULA) for the EventHubs container image. This is facilitated through the `WithAcceptEULA` function.
26-
26+
- [CosmosDB](#cosmosdb): `github.com/testcontainers/testcontainers-go/modules/azure/cosmosdb`.
2727
<!--codeinclude-->
2828
[Creating a Azurite container](../../modules/azure/azurite/examples_test.go) inside_block:runAzuriteContainer
2929
<!--/codeinclude-->
@@ -307,4 +307,49 @@ In the following example, inspired by the [Azure Event Hubs Go SDK](https://lear
307307
[Create Client](../../modules/azure/servicebus/examples_test.go) inside_block:createClient
308308
[Send messages to a Queue](../../modules/azure/servicebus/examples_test.go) inside_block:sendMessages
309309
[Receive messages from a Queue](../../modules/azure/servicebus/examples_test.go) inside_block:receiveMessages
310-
<!--/codeinclude-->
310+
<!--/codeinclude-->
311+
312+
## CosmosDB
313+
314+
### Run function
315+
316+
- Not available until the next release <a href="https://github.com/testcontainers/testcontainers-go"><span class="tc-version">:material-tag: main</span></a>
317+
318+
The CosmosDB module exposes one entrypoint function to create the CosmosDB container, and this function receives three parameters:
319+
320+
```golang
321+
func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustomizer) (*Container, error)
322+
```
323+
324+
- `context.Context`, the Go context.
325+
- `string`, the Docker image to use.
326+
- `testcontainers.ContainerCustomizer`, a variadic argument for passing options.
327+
328+
#### Image
329+
330+
Use the second argument in the `Run` function to set a valid Docker image.
331+
In example: `Run(context.Background(), "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:vnext-preview")`.
332+
333+
### Container Options
334+
335+
When starting the CosmosDB container, you can pass options in a variadic way to configure it.
336+
337+
{% include "../features/common_functional_options_list.md" %}
338+
339+
### Container Methods
340+
341+
The CosmosDB container exposes the following methods:
342+
343+
#### ConnectionString
344+
345+
- Not available until the next release <a href="https://github.com/testcontainers/testcontainers-go"><span class="tc-version">:material-tag: main</span></a>
346+
347+
Returns the connection string to connect to the CosmosDB container and an error, passing the Go context as parameter.
348+
349+
### Examples
350+
351+
#### Connect and Create database
352+
353+
<!--codeinclude-->
354+
[Connect_CreateDatabase](../../modules/azure/cosmosdb/examples_test.go) inside_block:ExampleRun_connect
355+
<!--/codeinclude-->

modules/azure/cosmosdb/cosmosdb.go

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
package cosmosdb
2+
3+
import (
4+
"context"
5+
"fmt"
6+
7+
"github.com/docker/go-connections/nat"
8+
9+
"github.com/testcontainers/testcontainers-go"
10+
"github.com/testcontainers/testcontainers-go/wait"
11+
)
12+
13+
const (
14+
defaultPort = "8081/tcp"
15+
defaultProtocol = "http"
16+
17+
// Well-known, publicly documented account key for the Azure CosmosDB Emulator.
18+
// See: https://learn.microsoft.com/en-us/azure/cosmos-db/how-to-develop-emulator
19+
testAccKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="
20+
)
21+
22+
// Container represents the CosmosDB container type used in the module
23+
type Container struct {
24+
testcontainers.Container
25+
}
26+
27+
// Run creates an instance of the CosmosDB container type
28+
func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustomizer) (*Container, error) {
29+
// Initialize with module defaults
30+
moduleOpts := []testcontainers.ContainerCustomizer{
31+
testcontainers.WithExposedPorts(defaultPort),
32+
testcontainers.WithCmdArgs("--enable-explorer", "false"),
33+
testcontainers.WithWaitStrategy(
34+
wait.ForAll(
35+
wait.ForLog("Started"),
36+
wait.ForListeningPort(nat.Port(defaultPort)),
37+
),
38+
),
39+
}
40+
41+
// Add user-provided options
42+
moduleOpts = append(moduleOpts, opts...)
43+
44+
ctr, err := testcontainers.Run(ctx, img, moduleOpts...)
45+
var c *Container
46+
if ctr != nil {
47+
c = &Container{Container: ctr}
48+
}
49+
50+
if err != nil {
51+
return c, fmt.Errorf("run cosmosdb: %w", err)
52+
}
53+
54+
return c, nil
55+
}
56+
57+
// ConnectionString returns a connection string that can be used to connect to the CosmosDB emulator.
58+
// The connection string includes the account endpoint (host:port) and the default test account key.
59+
// It returns an error if the port endpoint cannot be determined.
60+
//
61+
// Format: "AccountEndpoint=<host>:<port>;AccountKey=<accountKey>"
62+
func (c *Container) ConnectionString(ctx context.Context) (string, error) {
63+
endpoint, err := c.PortEndpoint(ctx, defaultPort, defaultProtocol)
64+
if err != nil {
65+
return "", fmt.Errorf("port endpoint: %w", err)
66+
}
67+
68+
return fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s;", endpoint, testAccKey), nil
69+
}
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
package cosmosdb_test
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"testing"
7+
8+
"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
9+
"github.com/stretchr/testify/require"
10+
11+
"github.com/testcontainers/testcontainers-go"
12+
"github.com/testcontainers/testcontainers-go/modules/azure/cosmosdb"
13+
)
14+
15+
func TestCosmosDB(t *testing.T) {
16+
ctx := context.Background()
17+
18+
ctr, err := cosmosdb.Run(ctx, "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:vnext-preview")
19+
testcontainers.CleanupContainer(t, ctr)
20+
require.NoError(t, err)
21+
22+
// Create Azure Cosmos client
23+
connStr, err := ctr.ConnectionString(ctx)
24+
require.NoError(t, err)
25+
require.NotNil(t, connStr)
26+
27+
p, err := cosmosdb.NewContainerPolicy(ctx, ctr)
28+
require.NoError(t, err)
29+
30+
client, err := azcosmos.NewClientFromConnectionString(connStr, p.ClientOptions())
31+
require.NoError(t, err)
32+
require.NotNil(t, client)
33+
34+
// Create database
35+
createDatabaseResp, err := client.CreateDatabase(ctx, azcosmos.DatabaseProperties{ID: "myDatabase"}, nil)
36+
require.NoError(t, err)
37+
require.NotNil(t, createDatabaseResp)
38+
39+
dbClient, err := client.NewDatabase("myDatabase")
40+
require.NoError(t, err)
41+
require.NotNil(t, dbClient)
42+
43+
// Create container
44+
containerProps := azcosmos.ContainerProperties{
45+
ID: "myContainer",
46+
PartitionKeyDefinition: azcosmos.PartitionKeyDefinition{Paths: []string{"/category"}},
47+
}
48+
createContainerResp, err := dbClient.CreateContainer(ctx, containerProps, nil)
49+
require.NoError(t, err)
50+
require.NotNil(t, createContainerResp)
51+
containerClient, err := dbClient.NewContainer("myContainer")
52+
require.NoError(t, err)
53+
require.NotNil(t, containerClient)
54+
55+
// Create item
56+
type Product struct {
57+
ID string `json:"id"`
58+
Category string `json:"category"`
59+
Name string `json:"name"`
60+
}
61+
62+
testItem := Product{ID: "item123", Category: "gear-surf-surfboards", Name: "Yamba Surfboard"}
63+
64+
pk := azcosmos.NewPartitionKeyString(testItem.Category)
65+
66+
jsonItem, err := json.Marshal(testItem)
67+
require.NoError(t, err)
68+
69+
createItemResp, err := containerClient.CreateItem(ctx, pk, jsonItem, nil)
70+
require.NoError(t, err)
71+
require.NotNil(t, createItemResp)
72+
73+
// Read item
74+
readItemResp, err := containerClient.ReadItem(ctx, pk, testItem.ID, nil)
75+
require.NoError(t, err)
76+
require.NotNil(t, readItemResp)
77+
}
Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
package cosmosdb_test
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"log"
7+
"net/http"
8+
9+
"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
10+
11+
"github.com/testcontainers/testcontainers-go"
12+
"github.com/testcontainers/testcontainers-go/modules/azure/cosmosdb"
13+
)
14+
15+
func ExampleRun() {
16+
ctx := context.Background()
17+
18+
cosmosdbContainer, err := cosmosdb.Run(ctx, "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:vnext-preview")
19+
defer func() {
20+
if err := testcontainers.TerminateContainer(cosmosdbContainer); err != nil {
21+
log.Printf("failed to terminate container: %s", err)
22+
}
23+
}()
24+
if err != nil {
25+
log.Printf("failed to start container: %s", err)
26+
return
27+
}
28+
// }
29+
30+
state, err := cosmosdbContainer.State(ctx)
31+
if err != nil {
32+
log.Printf("failed to get container state: %s", err)
33+
return
34+
}
35+
36+
fmt.Println(state.Running)
37+
38+
// Output:
39+
// true
40+
}
41+
42+
func ExampleRun_connect() {
43+
ctx := context.Background()
44+
45+
cosmosdbContainer, err := cosmosdb.Run(ctx, "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:vnext-preview")
46+
defer func() {
47+
if err := testcontainers.TerminateContainer(cosmosdbContainer); err != nil {
48+
log.Printf("failed to terminate container: %s", err)
49+
}
50+
}()
51+
if err != nil {
52+
log.Printf("failed to start container: %s", err)
53+
return
54+
}
55+
56+
connString, err := cosmosdbContainer.ConnectionString(ctx)
57+
if err != nil {
58+
log.Printf("failed to get connection string: %s", err)
59+
return
60+
}
61+
62+
p, err := cosmosdb.NewContainerPolicy(ctx, cosmosdbContainer)
63+
if err != nil {
64+
log.Printf("failed to create policy: %s", err)
65+
return
66+
}
67+
68+
client, err := azcosmos.NewClientFromConnectionString(connString, p.ClientOptions())
69+
if err != nil {
70+
log.Printf("failed to create cosmosdb client: %s", err)
71+
return
72+
}
73+
74+
createDatabaseResp, err := client.CreateDatabase(ctx, azcosmos.DatabaseProperties{ID: "myDatabase"}, nil)
75+
if err != nil {
76+
log.Printf("failed to create database: %s", err)
77+
return
78+
}
79+
// }
80+
81+
fmt.Println(createDatabaseResp.RawResponse.StatusCode == http.StatusCreated)
82+
83+
// Output:
84+
// true
85+
}

modules/azure/cosmosdb/policy.go

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
package cosmosdb
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"net/http"
7+
8+
"github.com/Azure/azure-sdk-for-go/sdk/azcore"
9+
"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
10+
"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
11+
)
12+
13+
// ContainerPolicy ensures that requests always target the CosmosDB emulator container endpoint.
14+
// It overrides the CosmosDB client's globalEndpointManager, which would otherwise dynamically
15+
// update [http.Request.Host] based on global endpoint discovery, pinning all requests to the container.
16+
type ContainerPolicy struct {
17+
endpoint string
18+
}
19+
20+
func NewContainerPolicy(ctx context.Context, c *Container) (*ContainerPolicy, error) {
21+
endpoint, err := c.PortEndpoint(ctx, defaultPort, "")
22+
if err != nil {
23+
return nil, fmt.Errorf("port endpoint: %w", err)
24+
}
25+
26+
return &ContainerPolicy{
27+
endpoint: endpoint,
28+
}, nil
29+
}
30+
31+
func (p *ContainerPolicy) Do(req *policy.Request) (*http.Response, error) {
32+
req.Raw().Host = p.endpoint
33+
req.Raw().URL.Host = p.endpoint
34+
35+
return req.Next()
36+
}
37+
38+
// ClientOptions returns Azure CosmosDB client options that contain ContainerPolicy.
39+
func (p *ContainerPolicy) ClientOptions() *azcosmos.ClientOptions {
40+
return &azcosmos.ClientOptions{
41+
ClientOptions: azcore.ClientOptions{
42+
PerRetryPolicies: []policy.Policy{p},
43+
},
44+
}
45+
}

modules/azure/go.mod

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ go 1.24.0
55
toolchain go1.24.7
66

77
require (
8-
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0
8+
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.19.1
9+
github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos v1.4.1
910
github.com/Azure/azure-sdk-for-go/sdk/data/aztables v1.3.0
1011
github.com/Azure/azure-sdk-for-go/sdk/messaging/azeventhubs v1.3.0
1112
github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus v1.8.0
@@ -19,8 +20,8 @@ require (
1920

2021
require (
2122
dario.cat/mergo v1.0.2 // indirect
22-
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.8.2 // indirect
23-
github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0 // indirect
23+
github.com/Azure/azure-sdk-for-go v68.0.0+incompatible // indirect
24+
github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.2 // indirect
2425
github.com/Azure/go-amqp v1.3.0 // indirect
2526
github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
2627
github.com/Microsoft/go-winio v0.6.2 // indirect

0 commit comments

Comments
 (0)