feature(api): add migration to insert vpn default value

Author: henrybarreto

api/store/mongo/migrations/main.go

@@ -84,6 +84,7 @@ func GenerateMigrations() []migrate.Migration {
 		migration72,
 		migration73,
 		migration74,
+		migration75,
 	}
 }
 

api/store/mongo/migrations/migration_75.go

@@ -0,0 +1,64 @@
+package migrations
+
+import (
+	"context"
+
+	"github.com/shellhub-io/shellhub/pkg/envs"
+	"github.com/sirupsen/logrus"
+	migrate "github.com/xakep666/mongo-migrate"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+)
+
+var migration75 = migrate.Migration{
+	Version:     75,
+	Description: "Adding VPN settings to namespace",
+	Up: migrate.MigrationFunc(func(ctx context.Context, db *mongo.Database) error {
+		logrus.WithFields(logrus.Fields{
+			"component": "migration",
+			"version":   75,
+			"action":    "Up",
+		}).Info("Applying migration")
+
+		if envs.IsEnterprise() {
+			update := bson.M{
+				"$set": bson.M{
+					"vpn": bson.M{
+						"enable":  false,
+						"address": bson.A{10, 0, 0, 0},
+						"mask":    16,
+					},
+				},
+			}
+
+			_, err := db.
+				Collection("namespaces").
+				UpdateMany(ctx, bson.M{}, update)
+
+			return err
+		}
+
+		return nil
+	}),
+	Down: migrate.MigrationFunc(func(ctx context.Context, db *mongo.Database) error {
+		logrus.WithFields(logrus.Fields{
+			"component": "migration",
+			"version":   75,
+			"action":    "Down",
+		}).Info("Reverting migration")
+
+		if envs.IsEnterprise() {
+			update := bson.M{
+				"$unset": bson.M{"vpn": ""},
+			}
+
+			_, err := db.
+				Collection("namespaces").
+				UpdateMany(ctx, bson.M{}, update)
+
+			return err
+		}
+
+		return nil
+	}),
+}

api/store/mongo/migrations/migration_75_test.go

@@ -0,0 +1,123 @@
+package migrations
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/shellhub-io/shellhub/pkg/envs"
+	env_mocks "github.com/shellhub-io/shellhub/pkg/envs/mocks"
+	"github.com/shellhub-io/shellhub/pkg/models"
+	"github.com/stretchr/testify/assert"
+	migrate "github.com/xakep666/mongo-migrate"
+	"go.mongodb.org/mongo-driver/bson"
+)
+
+func TestMigration75(t *testing.T) {
+	ctx := context.Background()
+
+	envMock = &env_mocks.Backend{}
+	envs.DefaultBackend = envMock
+
+	cases := []struct {
+		description  string
+		setup        func() error
+		requireMocks func()
+		test         func() error
+	}{
+		{
+			description: "Success to apply up on migration 75",
+			setup: func() error {
+				_, err := c.
+					Database("test").
+					Collection("namespaces").
+					InsertOne(ctx, models.Namespace{
+						TenantID: "00000000-0000-4000-0000-000000000000",
+					})
+
+				return err
+			},
+			requireMocks: func() {
+				envMock.On("Get", "SHELLHUB_ENTERPRISE").Return("true").Once()
+			},
+			test: func() error {
+				migrations := GenerateMigrations()[74:75]
+				migrates := migrate.NewMigrate(c.Database("test"), migrations...)
+				err := migrates.Up(context.Background(), migrate.AllAvailable)
+				if err != nil {
+					return err
+				}
+
+				query := c.
+					Database("test").
+					Collection("namespaces").
+					FindOne(context.TODO(), bson.M{"tenant_id": "00000000-0000-4000-0000-000000000000"})
+
+				ns := new(models.Namespace)
+				if err := query.Decode(ns); err != nil {
+					return errors.New("unable to find the namespace")
+				}
+
+				if ns.VPN == nil {
+					return errors.New("unable to apply the migration")
+				}
+
+				return nil
+			},
+		}, {
+			description: "Success to unapply the migration 75",
+			setup: func() error {
+				_, err := c.
+					Database("test").
+					Collection("namespaces").
+					InsertOne(ctx, models.Namespace{
+						TenantID: "00000000-0000-4000-0000-000000000000",
+						Settings: &models.NamespaceSettings{},
+					})
+
+				return err
+			},
+			requireMocks: func() {},
+			test: func() error {
+				migrations := GenerateMigrations()[74:75]
+				migrates := migrate.NewMigrate(c.Database("test"), migrations...)
+				err := migrates.Down(context.Background(), migrate.AllAvailable)
+				if err != nil {
+					return err
+				}
+
+				query := c.
+					Database("test").
+					Collection("namespaces").
+					FindOne(context.TODO(), bson.M{"tenant_id": "00000000-0000-4000-0000-000000000000"})
+
+				ns := new(models.Namespace)
+				if err := query.Decode(ns); err != nil {
+					return errors.New("unable to find the namespace")
+				}
+
+				if ns.VPN != nil {
+					return errors.New("unable to unapply the migration")
+				}
+
+				return nil
+			},
+		},
+	}
+
+	for _, test := range cases {
+		tc := test
+		t.Run(tc.description, func(t *testing.T) {
+			tc.requireMocks()
+
+			t.Cleanup(func() {
+				assert.NoError(t, srv.Reset())
+			})
+
+			assert.NoError(t, tc.setup())
+			assert.NoError(t, tc.test())
+
+			envMock.AssertExpectations(t)
+		})
+	}
+}