crawlab/core/models/client/model_service_test.go

package client_test

import (
	"context"
	"testing"
	"time"

	"github.com/crawlab-team/crawlab/core/models/models/v2"

	"github.com/apex/log"
	"github.com/crawlab-team/crawlab/core/grpc/server"
	"github.com/crawlab-team/crawlab/core/models/client"
	"github.com/crawlab-team/crawlab/core/models/service"
	"github.com/crawlab-team/crawlab/db/mongo"
	"github.com/spf13/viper"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.mongodb.org/mongo-driver/bson"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
)

func setupTestDB() {
	viper.Set("mongo.db", "testdb")
}

func teardownTestDB() {
	db := mongo.GetMongoDb("testdb")
	db.Drop(context.Background())
}

func startSvr(svr *server.GrpcServer) {
	err := svr.Start()
	if err != nil {
		log.Errorf("failed to start grpc server: %v", err)
	}
}

func stopSvr(svr *server.GrpcServer) {
	err := svr.Stop()
	if err != nil {
		log.Errorf("failed to stop grpc server: %v", err)
	}
}

func TestModelService_GetById(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	id, err := modelSvc.InsertOne(m)
	require.Nil(t, err)
	m.SetId(id)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	res, err := clientSvc.GetById(m.Id)
	require.Nil(t, err)
	assert.Equal(t, res.Id, m.Id)
	assert.Equal(t, res.Name, m.Name)
}

func TestModelService_GetOne(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	id, err := modelSvc.InsertOne(m)
	require.Nil(t, err)
	m.SetId(id)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	res, err := clientSvc.GetOne(bson.M{"name": m.Name}, nil)
	require.Nil(t, err)
	assert.Equal(t, res.Id, m.Id)
	assert.Equal(t, res.Name, m.Name)
}

func TestModelService_GetMany(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	id, err := modelSvc.InsertOne(m)
	require.Nil(t, err)
	m.SetId(id)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	res, err := clientSvc.GetMany(bson.M{"name": m.Name}, nil)
	require.Nil(t, err)
	assert.Equal(t, len(res), 1)
	assert.Equal(t, res[0].Id, m.Id)
	assert.Equal(t, res[0].Name, m.Name)
}

func TestModelService_DeleteById(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	id, err := modelSvc.InsertOne(m)
	require.Nil(t, err)
	m.SetId(id)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	err = clientSvc.DeleteById(m.Id)
	require.Nil(t, err)

	res, err := clientSvc.GetById(m.Id)
	assert.NotNil(t, err)
	require.Nil(t, res)
}

func TestModelService_DeleteOne(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	id, err := modelSvc.InsertOne(m)
	require.Nil(t, err)
	m.SetId(id)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	err = clientSvc.DeleteOne(bson.M{"name": m.Name})
	require.Nil(t, err)

	res, err := clientSvc.GetOne(bson.M{"name": m.Name}, nil)
	assert.NotNil(t, err)
	require.Nil(t, res)
}

func TestModelService_DeleteMany(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	id, err := modelSvc.InsertOne(m)
	require.Nil(t, err)
	m.SetId(id)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	err = clientSvc.DeleteMany(bson.M{"name": m.Name})
	require.Nil(t, err)

	res, err := clientSvc.GetMany(bson.M{"name": m.Name}, nil)
	require.Nil(t, err)
	assert.Equal(t, len(res), 0)
}

func TestModelService_UpdateById(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	id, err := modelSvc.InsertOne(m)
	require.Nil(t, err)
	m.SetId(id)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	err = clientSvc.UpdateById(m.Id, bson.M{"$set": bson.M{"name": "New Name"}})
	require.Nil(t, err)

	res, err := clientSvc.GetById(m.Id)
	require.Nil(t, err)
	assert.Equal(t, res.Name, "New Name")
}

func TestModelService_UpdateOne(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	id, err := modelSvc.InsertOne(m)
	require.Nil(t, err)
	m.SetId(id)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	err = clientSvc.UpdateOne(bson.M{"name": m.Name}, bson.M{"$set": bson.M{"name": "New Name"}})
	require.Nil(t, err)

	res, err := clientSvc.GetOne(bson.M{"name": "New Name"}, nil)
	require.Nil(t, err)
	assert.Equal(t, res.Name, "New Name")
}

func TestModelService_UpdateMany(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m1 := models.TestModelV2{
		Name: "Test Name",
	}
	m2 := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	_, err = modelSvc.InsertOne(m1)
	require.Nil(t, err)
	_, err = modelSvc.InsertOne(m2)
	require.Nil(t, err)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	err = clientSvc.UpdateMany(bson.M{"name": "Test Name"}, bson.M{"$set": bson.M{"name": "New Name"}})
	require.Nil(t, err)

	res, err := clientSvc.GetMany(bson.M{"name": "New Name"}, nil)
	require.Nil(t, err)
	assert.Equal(t, len(res), 2)
}

func TestModelService_ReplaceById(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	id, err := modelSvc.InsertOne(m)
	require.Nil(t, err)
	m.SetId(id)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	m.Name = "New Name"
	err = clientSvc.ReplaceById(m.Id, m)
	require.Nil(t, err)

	res, err := clientSvc.GetById(m.Id)
	require.Nil(t, err)
	assert.Equal(t, res.Name, "New Name")
}

func TestModelService_ReplaceOne(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	m := models.TestModelV2{
		Name: "Test Name",
	}
	modelSvc := service.NewModelService[models.TestModelV2]()
	id, err := modelSvc.InsertOne(m)
	require.Nil(t, err)
	m.SetId(id)
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	m.Name = "New Name"
	err = clientSvc.ReplaceOne(bson.M{"name": "Test Name"}, m)
	require.Nil(t, err)

	res, err := clientSvc.GetOne(bson.M{"name": "New Name"}, nil)
	require.Nil(t, err)
	assert.Equal(t, res.Name, "New Name")
}

func TestModelService_InsertOne(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	m := models.TestModelV2{
		Name: "Test Name",
	}
	id, err := clientSvc.InsertOne(m)
	require.Nil(t, err)

	res, err := clientSvc.GetById(id)
	require.Nil(t, err)
	assert.Equal(t, res.Name, m.Name)
}

func TestModelService_InsertMany(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	testModels := []models.TestModelV2{
		{Name: "Test Name 1"},
		{Name: "Test Name 2"},
	}
	ids, err := clientSvc.InsertMany(testModels)
	require.Nil(t, err)

	for i, id := range ids {
		res, err := clientSvc.GetById(id)
		require.Nil(t, err)
		assert.Equal(t, res.Name, testModels[i].Name)
	}
}

func TestModelService_Count(t *testing.T) {
	setupTestDB()
	defer teardownTestDB()
	svr, err := server.NewGrpcServer()
	require.Nil(t, err)
	go startSvr(svr)
	defer stopSvr(svr)

	modelSvc := service.NewModelService[models.TestModelV2]()
	for i := 0; i < 5; i++ {
		_, err = modelSvc.InsertOne(models.TestModelV2{
			Name: "Test Name",
		})
		require.Nil(t, err)
	}
	time.Sleep(100 * time.Millisecond)

	c, err := grpc.NewClient("localhost:9666", grpc.WithTransportCredentials(insecure.NewCredentials()))
	require.Nil(t, err)
	c.Connect()

	clientSvc := client.NewModelService[models.TestModelV2]()
	count, err := clientSvc.Count(bson.M{})
	require.Nil(t, err)

	assert.Equal(t, count, 5)
}