feat: Add method for getting default database

kapetacom · Apr 16, 2024 · 11f69b6 · 11f69b6
1 parent 0ab4328
commit 11f69b6
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 47 deletions.
diff --git a/mongo.go b/mongo.go
@@ -18,10 +18,20 @@ const RESOURCE_PORT = "mongodb"
 
 type MongoDB struct {
 	*mongo.Client
+	dbName string
+}
+
+func (m *MongoDB) DB() *mongo.Database {
+	return m.Database(m.dbName)
 }
 
 func NewMongoDB(config providers.ConfigProvider, resourceName string) (*MongoDB, error) {
-	url, err := createConnectionString(config, resourceName)
+	resInfo, err := config.GetResourceInfo(RESOURCE_TYPE, RESOURCE_PORT, resourceName)
+	if err != nil {
+		return nil, err
+	}
+
+	url, err := createConnectionString(resInfo, resourceName)
 	if err != nil {
 		return nil, err
 	}
@@ -41,14 +51,10 @@ func NewMongoDB(config providers.ConfigProvider, resourceName string) (*MongoDB,
 	}
 	log.Printf("Connected successfully to mongodb database: %s\n", resourceName)
 
-	return &MongoDB{client}, nil
+	return &MongoDB{client, getDBName(resInfo, resourceName)}, nil
 }
 
-func createConnectionString(config providers.ConfigProvider, resourceName string) (string, error) {
-	resInfo, err := config.GetResourceInfo(RESOURCE_TYPE, RESOURCE_PORT, resourceName)
-	if err != nil {
-		return "", err
-	}
+func createConnectionString(resInfo *providers.ResourceInfo, resourceName string) (string, error) {
 	protocol := getProtocol(resInfo)
 	dbName := getDBName(resInfo, resourceName)
 

diff --git a/mongo_test.go b/mongo_test.go
@@ -3,7 +3,6 @@ package mongo
 import (
 	"testing"
 
-	config "github.com/kapetacom/sdk-go-config"
 	"github.com/kapetacom/sdk-go-config/providers"
 	"github.com/stretchr/testify/assert"
 )
@@ -50,12 +49,7 @@ func TestCreateConnectionStringProtocolMongodb(t *testing.T) {
 		dbName := "test"
 		urlWithPort := "mongodb://user:password@localhost:27017/test?authSource=admin&directConnection=true"
 
-		config := &config.ConfigProviderMock{
-			GetResourceInfoFunc: func(resourceType, resourcePort, resourceName string) (*providers.ResourceInfo, error) {
-				return resInfo, nil
-			},
-		}
-		actual, err := createConnectionString(config, dbName)
+		actual, err := createConnectionString(resInfo, dbName)
 		assert.NoError(t, err)
 		assert.Equal(t, urlWithPort, actual)
 	})
@@ -70,12 +64,7 @@ func TestCreateConnectionStringProtocolMongodb(t *testing.T) {
 		dbName := "test"
 		urlWithNoPort := "mongodb+srv://user:password@localhost/test?authSource=admin"
 
-		config := &config.ConfigProviderMock{
-			GetResourceInfoFunc: func(resourceType, resourcePort, resourceName string) (*providers.ResourceInfo, error) {
-				return resInfo, nil
-			},
-		}
-		actual, err := createConnectionString(config, dbName)
+		actual, err := createConnectionString(resInfo, dbName)
 		assert.NoError(t, err)
 		assert.Equal(t, urlWithNoPort, actual)
 	})
@@ -89,12 +78,7 @@ func TestCreateConnectionStringProtocolMongodb(t *testing.T) {
 		dbName := "test"
 		urlWithPort := "mongodb://user:password@localhost:27017/test?authSource=admin&directConnection=true"
 
-		config := &config.ConfigProviderMock{
-			GetResourceInfoFunc: func(resourceType, resourcePort, resourceName string) (*providers.ResourceInfo, error) {
-				return resInfo, nil
-			},
-		}
-		actual, err := createConnectionString(config, dbName)
+		actual, err := createConnectionString(resInfo, dbName)
 		assert.NoError(t, err)
 		assert.Equal(t, urlWithPort, actual)
 	})
@@ -109,12 +93,7 @@ func TestCreateConnectionStringProtocolMongodb(t *testing.T) {
 		dbName := "test"
 		urlWithPort := "mongodb://user:password@localhost:27017/test?authSource=admin&directConnection=true&ssl=true"
 
-		config := &config.ConfigProviderMock{
-			GetResourceInfoFunc: func(resourceType, resourcePort, resourceName string) (*providers.ResourceInfo, error) {
-				return resInfo, nil
-			},
-		}
-		actual, err := createConnectionString(config, dbName)
+		actual, err := createConnectionString(resInfo, dbName)
 		assert.NoError(t, err)
 		assert.Equal(t, urlWithPort, actual)
 	})
@@ -129,31 +108,22 @@ func TestCreateConnectionStringProtocolMongodb(t *testing.T) {
 		dbName := "test"
 		urlWithPort := "mongodb://user:password@localhost:27017/test?authSource=admin&directConnection=true&ssl=false"
 
-		config := &config.ConfigProviderMock{
-			GetResourceInfoFunc: func(resourceType, resourcePort, resourceName string) (*providers.ResourceInfo, error) {
-				return resInfo, nil
-			},
-		}
-		actual, err := createConnectionString(config, dbName)
+		actual, err := createConnectionString(resInfo, dbName)
 		assert.NoError(t, err)
 		assert.Equal(t, urlWithPort, actual)
 	})
 }
 
 func TestCreateConnectionString(t *testing.T) {
 	t.Run("create connection string", func(t *testing.T) {
-		config := &config.ConfigProviderMock{
-			GetResourceInfoFunc: func(resourceType, resourcePort, resourceName string) (*providers.ResourceInfo, error) {
-				return &providers.ResourceInfo{
-					Host:        "localhost",
-					Port:        "27017",
-					Credentials: map[string]string{"username": "user", "password": "password"},
-				}, nil
-			},
+		resInfo := &providers.ResourceInfo{
+			Host:        "localhost",
+			Port:        "27017",
+			Credentials: map[string]string{"username": "user", "password": "password"},
 		}
 		resourceName := "test"
 		expected := "mongodb://user:password@localhost:27017/test?authSource=admin&directConnection=true"
-		actual, err := createConnectionString(config, resourceName)
+		actual, err := createConnectionString(resInfo, resourceName)
 		assert.NoError(t, err)
 		assert.Equal(t, expected, actual)
 	})