<commit_before>package provisioner

import (
	"context"
	"io"
	"io/ioutil"
	"sync"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/mock"

	"github.com/supergiant/control/pkg/clouds"
	"github.com/supergiant/control/pkg/model"
	"github.com/supergiant/control/pkg/profile"
	"github.com/supergiant/control/pkg/sgerrors"
	"github.com/supergiant/control/pkg/testutils"
	"github.com/supergiant/control/pkg/workflows"
	"github.com/supergiant/control/pkg/workflows/steps"
)

type bufferCloser struct {
	io.Writer
	err error
}

func (b *bufferCloser) Close() error {
	return b.err
}

type mockKubeService struct {
	getError  error
	createErr error
	lock      sync.RWMutex
	data      map[string]*model.Kube
}

func (m *mockKubeService) Create(ctx context.Context, k *model.Kube) error {
	m.lock.Lock()
	m.data[k.ID] = k
	m.lock.Unlock()
	return m.createErr
}

func (m *mockKubeService) Get(ctx context.Context, kname string) (*model.Kube, error) {
	m.lock.RLock()
	defer m.lock.RUnlock()
	return m.data[kname], m.getError
}

type mockStep struct {
}

func (m *mockStep) Run(context.Context, io.Writer, *steps.Config) error {
	return nil
}

func (m *mockStep) Name() string {
	return ""
}

func (m *mockStep) Description() string {
	return ""
}

func (m *mockStep) Depends() []string {
	return nil
}

func (m *mockStep) Rollback(context.Context, io.Writer, *steps.Config) error {
	return nil
}

func TestNewProvisioner(t *testing.T) {
	storage := &testutils.MockStorage{}
	service := &mockKubeService{}
	interval := time.Second * 1

	p := NewProvisioner(storage, service, interval)

	if p.repository != storage {
		t.Errorf("Wrong repository expected %v actual %v",
			storage, p.repository)
	}

	if p.kubeService != service {
		t.Errorf("Wrong kube service value expected %v actual %v",
			service, p.kubeService)
	}

	if p.cancelMap == nil {
		t.Errorf("Cancel map must not be nil")
	}
}

func TestProvisionCluster(t *testing.T) {
	repository := &testutils.MockStorage{}
	repository.On("Put", mock.Anything,
		mock.Anything, mock.Anything,
		mock.Anything).Return(nil)

	bc := &bufferCloser{
		ioutil.Discard,
		nil,
	}

	svc := &mockKubeService{
		data: make(map[string]*model.Kube),
	}

	provisioner := TaskProvisioner{
		svc,
		repository,
		func(string) (io.WriteCloser, error) {
			return bc, nil
		},
		NewRateLimiter(time.Nanosecond * 1),
		make(map[string]func()),
	}

	workflows.Init()
	workflows.RegisterWorkFlow(workflows.ProvisionMaster, []steps.Step{
		&mockStep{},
	})
	workflows.RegisterWorkFlow(workflows.ProvisionNode, []steps.Step{
		&mockStep{},
	})
	workflows.RegisterWorkFlow(workflows.PostProvision, []steps.Step{
		&mockStep{},
	})
	workflows.RegisterWorkFlow(workflows.PreProvision, []steps.Step{
		&mockStep{},
	})

	p := &profile.Profile{
		Provider: clouds.DigitalOcean,
		MasterProfiles: []profile.NodeProfile{
			{
				"size":  "s-1vcpu-2gb",
				"image": "ubuntu-18-04-x64",
			},
		},
		NodesProfiles: []profile.NodeProfile{
			{
				"size":  "s-2vcpu-4gb",
				"image": "ubuntu-18-04-x64",
			},
			{
				"size":  "s-2vcpu-4gb",
				"image": "ubuntu-18-04-x64",
			},
		},
	}

	cfg, err := steps.NewConfig("", "", *p)

	if err != nil {
		t.Errorf("Unexpected error %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	taskMap, err := provisioner.ProvisionCluster(ctx, p, cfg)
	time.Sleep(time.Millisecond * 10)

	if err != nil {
		t.Errorf("Unexpected error %v while provisionCluster", err)
	}

	if len(taskMap) != 4 {
		t.Errorf("Expected task map len 4 actual %d", len(taskMap))
	}

	if len(taskMap["master"])+len(taskMap["node"]) + 1 != len(p.MasterProfiles)+len(p.NodesProfiles) {
		t.Errorf("Wrong task count expected %d actual %d",
			len(p.MasterProfiles)+
				len(p.NodesProfiles) + 1,
			len(taskMap))
	}

	if len(provisioner.cancelMap) != 1 {
		t.Errorf("Unexpected size of cancel map expected %d actual %d",
			1, len(provisioner.cancelMap))
	}

	// Cancel context to shut down cluster state monitoring
	cancel()
	if k, err := svc.Get(context.Background(), cfg.ClusterID); err == nil && k == nil {
		t.Errorf("Kube %s not found", k.ID)

		if len(k.Tasks) != len(p.MasterProfiles)+len(p.NodesProfiles)+1 {
			t.Errorf("Wrong count of tasks in kube expected %d actual %d",
				len(p.MasterProfiles)+len(p.NodesProfiles)+1, len(k.Tasks))
		}
	}
}

func TestProvisionNodes(t *testing.T) {
	repository := &testutils.MockStorage{}<commit_msg>Fix CI tests (#1358)<commit_after>package provisioner

import (
	"context"
	"io"
	"io/ioutil"
	"sync"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/mock"

	"github.com/supergiant/control/pkg/clouds"
	"github.com/supergiant/control/pkg/model"
	"github.com/supergiant/control/pkg/profile"
	"github.com/supergiant/control/pkg/sgerrors"
	"github.com/supergiant/control/pkg/testutils"
	"github.com/supergiant/control/pkg/workflows"
	"github.com/supergiant/control/pkg/workflows/steps"
)

type bufferCloser struct {
	io.Writer
	err error
}

func (b *bufferCloser) Close() error {
	return b.err
}

type mockKubeService struct {
	getError  error
	createErr error
	lock      sync.RWMutex
	data      map[string]*model.Kube
}

func (m *mockKubeService) Create(ctx context.Context, k *model.Kube) error {
	m.lock.Lock()
	m.data[k.ID] = k
	m.lock.Unlock()
	return m.createErr
}

func (m *mockKubeService) Get(ctx context.Context, kname string) (*model.Kube, error) {
	m.lock.RLock()
	defer m.lock.RUnlock()
	return m.data[kname], m.getError
}

type mockStep struct {
}

func (m *mockStep) Run(context.Context, io.Writer, *steps.Config) error {
	return nil
}

func (m *mockStep) Name() string {
	return ""
}

func (m *mockStep) Description() string {
	return ""
}

func (m *mockStep) Depends() []string {
	return nil
}

func (m *mockStep) Rollback(context.Context, io.Writer, *steps.Config) error {
	return nil
}

func TestNewProvisioner(t *testing.T) {
	storage := &testutils.MockStorage{}
	service := &mockKubeService{}
	interval := time.Second * 1

	p := NewProvisioner(storage, service, interval)

	if p.repository != storage {
		t.Errorf("Wrong repository expected %v actual %v",
			storage, p.repository)
	}

	if p.kubeService != service {
		t.Errorf("Wrong kube service value expected %v actual %v",
			service, p.kubeService)
	}

	if p.cancelMap == nil {
		t.Errorf("Cancel map must not be nil")
	}
}

func TestProvisionCluster(t *testing.T) {
	repository := &testutils.MockStorage{}
	repository.On("Put", mock.Anything,
		mock.Anything, mock.Anything,
		mock.Anything).Return(nil)

	bc := &bufferCloser{
		ioutil.Discard,
		nil,
	}

	svc := &mockKubeService{
		data: make(map[string]*model.Kube),
	}

	provisioner := TaskProvisioner{
		svc,
		repository,
		func(string) (io.WriteCloser, error) {
			return bc, nil
		},
		NewRateLimiter(time.Nanosecond * 1),
		make(map[string]func()),
	}

	workflows.Init()
	workflows.RegisterWorkFlow(workflows.ProvisionMaster, []steps.Step{
		&mockStep{},
	})
	workflows.RegisterWorkFlow(workflows.ProvisionNode, []steps.Step{
		&mockStep{},
	})
	workflows.RegisterWorkFlow(workflows.PostProvision, []steps.Step{
		&mockStep{},
	})
	workflows.RegisterWorkFlow(workflows.PreProvision, []steps.Step{
		&mockStep{},
	})

	p := &profile.Profile{
		Provider: clouds.DigitalOcean,
		MasterProfiles: []profile.NodeProfile{
			{
				"size":  "s-1vcpu-2gb",
				"image": "ubuntu-18-04-x64",
			},
		},
		NodesProfiles: []profile.NodeProfile{
			{
				"size":  "s-2vcpu-4gb",
				"image": "ubuntu-18-04-x64",
			},
			{
				"size":  "s-2vcpu-4gb",
				"image": "ubuntu-18-04-x64",
			},
		},
	}

	cfg, err := steps.NewConfig("", "", *p)

	if err != nil {
		t.Errorf("Unexpected error %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	taskMap, err := provisioner.ProvisionCluster(ctx, p, cfg)
	time.Sleep(time.Millisecond * 10)

	if err != nil {
		t.Errorf("Unexpected error %v while provisionCluster", err)
	}

	if len(taskMap) != 4 {
		t.Errorf("Expected task map len 4 actual %d", len(taskMap))
	}

	if len(taskMap["master"])+len(taskMap["node"]) != len(p.MasterProfiles)+len(p.NodesProfiles) {
		t.Errorf("Wrong task count expected %d actual %d",
			len(p.MasterProfiles)+
				len(p.NodesProfiles),
			len(taskMap["master"])+len(taskMap["node"]))
	}

	if len(provisioner.cancelMap) != 1 {
		t.Errorf("Unexpected size of cancel map expected %d actual %d",
			1, len(provisioner.cancelMap))
	}

	// Cancel context to shut down cluster state monitoring
	cancel()
	if k, err := svc.Get(context.Background(), cfg.ClusterID); err == nil && k == nil {
		t.Errorf("Kube %s not found", k.ID)

		if len(k.Tasks) != len(p.MasterProfiles)+len(p.NodesProfiles)+1 {
			t.Errorf("Wrong count of tasks in kube expected %d actual %d",
				len(p.MasterProfiles)+len(p.NodesProfiles)+1, len(k.Tasks))
		}
	}
}

func TestProvisionNodes(t *testing.T) {
	repository := &testutils.MockStorage{}
{'fraction_non_alphanumeric': 0.1600896860986547, 'fraction_numerical': 0.008295964125560539, 'mean_word_length': 4.156647398843931, 'pattern_counts': {'":': 12, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7539519', 'n_tokens_mistral': 3625, 'n_tokens_neox': 3361, 'n_words': 672}