<commit_before>package client

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"text/template"
	"time"

	"github.com/boltdb/bolt"
	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/nomad/command/agent/consul"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/helper/uuid"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"
	"github.com/hashicorp/nomad/version"
	"github.com/kr/pretty"
	"github.com/stretchr/testify/assert"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/vaultclient"
	"github.com/stretchr/testify/require"
)

type MockAllocStateUpdater struct {
	Allocs []*structs.Allocation
	mu     sync.Mutex
}

// Update fulfills the TaskStateUpdater interface
func (m *MockAllocStateUpdater) Update(alloc *structs.Allocation) {
	m.mu.Lock()
	m.Allocs = append(m.Allocs, alloc)
	m.mu.Unlock()
}

// Last returns a copy of the last alloc (or nil) sync'd
func (m *MockAllocStateUpdater) Last() *structs.Allocation {
	m.mu.Lock()
	defer m.mu.Unlock()
	n := len(m.Allocs)
	if n == 0 {
		return nil
	}
	return m.Allocs[n-1].Copy()
}

// allocationBucketExists checks if the allocation bucket was created.
func allocationBucketExists(tx *bolt.Tx, allocID string) bool {
	allocations := tx.Bucket(allocationsBucket)
	if allocations == nil {
		return false
	}

	// Retrieve the specific allocations bucket
	alloc := allocations.Bucket([]byte(allocID))
	return alloc != nil
}

func testAllocRunnerFromAlloc(t *testing.T, alloc *structs.Allocation, restarts bool) (*MockAllocStateUpdater, *AllocRunner) {
	conf := config.DefaultConfig()
	conf.Node = mock.Node()
	conf.StateDir = os.TempDir()
	conf.AllocDir = os.TempDir()
	tmp, _ := ioutil.TempFile("", "state-db")
	db, _ := bolt.Open(tmp.Name(), 0600, nil)
	upd := &MockAllocStateUpdater{}
	if !restarts {
		*alloc.Job.LookupTaskGroup(alloc.TaskGroup).RestartPolicy = structs.RestartPolicy{Attempts: 0}
		alloc.Job.Type = structs.JobTypeBatch
	}
	vclient := vaultclient.NewMockVaultClient()
	ar := NewAllocRunner(testlog.Logger(t), conf, db, upd.Update, alloc, vclient, newMockConsulServiceClient(t), noopPrevAlloc{})
	return upd, ar
}

func testAllocRunner(t *testing.T, restarts bool) (*MockAllocStateUpdater, *AllocRunner) {
	// Use mock driver
	alloc := mock.Alloc()
	task := alloc.Job.TaskGroups[0].Tasks[0]
	task.Driver = "mock_driver"
	task.Config["run_for"] = "500ms"
	return testAllocRunnerFromAlloc(t, alloc, restarts)
}

func TestAllocRunner_SimpleRun(t *testing.T) {
	t.Parallel()
	upd, ar := testAllocRunner(t, false)
	go ar.Run()
	defer ar.Destroy()

	testutil.WaitForResult(func() (bool, error) {
		last := upd.Last()
		if last == nil {
			return false, fmt.Errorf("No updates")
		}
		if last.ClientStatus != structs.AllocClientStatusComplete {
			return false, fmt.Errorf("got status %v; want %v", last.ClientStatus, structs.AllocClientStatusComplete)
		}
		return true, nil
	}, func(err error) {
		t.Fatalf("err: %v", err)
	})
}

// Test that FinisheAt is set when the alloc is in a terminal state
func TestAllocRunner_FinishedAtSet(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, ar := testAllocRunner(t, false)
	//ar.alloc.ClientStatus = structs.AllocClientStatusFailed
	ar.allocClientStatus = structs.AllocClientStatusFailed
	alloc := ar.Alloc()
	taskFinishedAt := make(map[string]time.Time)
	require.NotEmpty(alloc.TaskStates)
	for name, s := range alloc.TaskStates {
		require.False(s.FinishedAt.IsZero())
		taskFinishedAt[name] = s.FinishedAt
	}

	// Verify that calling again should not mutate finishedAt
	alloc2 := ar.Alloc()
	for name, s := range alloc2.TaskStates {
		require.Equal(taskFinishedAt[name], s.FinishedAt)
	}

}

// Test that FinisheAt is set when the alloc is in a terminal state
func TestAllocRunner_FinishedAtSet_TaskEvents(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, ar := testAllocRunner(t, false)
	ar.taskStates[ar.alloc.Job.TaskGroups[0].Tasks[0].Name] = &structs.TaskState{State: structs.TaskStateDead, Failed: true}

	alloc := ar.Alloc()
	taskFinishedAt := make(map[string]time.Time)
	require.NotEmpty(alloc.TaskStates)
	for name, s := range alloc.TaskStates {<commit_msg>remove outdated commented out test code<commit_after>package client

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"text/template"
	"time"

	"github.com/boltdb/bolt"
	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/nomad/command/agent/consul"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/helper/uuid"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"
	"github.com/hashicorp/nomad/version"
	"github.com/kr/pretty"
	"github.com/stretchr/testify/assert"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/vaultclient"
	"github.com/stretchr/testify/require"
)

type MockAllocStateUpdater struct {
	Allocs []*structs.Allocation
	mu     sync.Mutex
}

// Update fulfills the TaskStateUpdater interface
func (m *MockAllocStateUpdater) Update(alloc *structs.Allocation) {
	m.mu.Lock()
	m.Allocs = append(m.Allocs, alloc)
	m.mu.Unlock()
}

// Last returns a copy of the last alloc (or nil) sync'd
func (m *MockAllocStateUpdater) Last() *structs.Allocation {
	m.mu.Lock()
	defer m.mu.Unlock()
	n := len(m.Allocs)
	if n == 0 {
		return nil
	}
	return m.Allocs[n-1].Copy()
}

// allocationBucketExists checks if the allocation bucket was created.
func allocationBucketExists(tx *bolt.Tx, allocID string) bool {
	allocations := tx.Bucket(allocationsBucket)
	if allocations == nil {
		return false
	}

	// Retrieve the specific allocations bucket
	alloc := allocations.Bucket([]byte(allocID))
	return alloc != nil
}

func testAllocRunnerFromAlloc(t *testing.T, alloc *structs.Allocation, restarts bool) (*MockAllocStateUpdater, *AllocRunner) {
	conf := config.DefaultConfig()
	conf.Node = mock.Node()
	conf.StateDir = os.TempDir()
	conf.AllocDir = os.TempDir()
	tmp, _ := ioutil.TempFile("", "state-db")
	db, _ := bolt.Open(tmp.Name(), 0600, nil)
	upd := &MockAllocStateUpdater{}
	if !restarts {
		*alloc.Job.LookupTaskGroup(alloc.TaskGroup).RestartPolicy = structs.RestartPolicy{Attempts: 0}
		alloc.Job.Type = structs.JobTypeBatch
	}
	vclient := vaultclient.NewMockVaultClient()
	ar := NewAllocRunner(testlog.Logger(t), conf, db, upd.Update, alloc, vclient, newMockConsulServiceClient(t), noopPrevAlloc{})
	return upd, ar
}

func testAllocRunner(t *testing.T, restarts bool) (*MockAllocStateUpdater, *AllocRunner) {
	// Use mock driver
	alloc := mock.Alloc()
	task := alloc.Job.TaskGroups[0].Tasks[0]
	task.Driver = "mock_driver"
	task.Config["run_for"] = "500ms"
	return testAllocRunnerFromAlloc(t, alloc, restarts)
}

func TestAllocRunner_SimpleRun(t *testing.T) {
	t.Parallel()
	upd, ar := testAllocRunner(t, false)
	go ar.Run()
	defer ar.Destroy()

	testutil.WaitForResult(func() (bool, error) {
		last := upd.Last()
		if last == nil {
			return false, fmt.Errorf("No updates")
		}
		if last.ClientStatus != structs.AllocClientStatusComplete {
			return false, fmt.Errorf("got status %v; want %v", last.ClientStatus, structs.AllocClientStatusComplete)
		}
		return true, nil
	}, func(err error) {
		t.Fatalf("err: %v", err)
	})
}

// Test that FinisheAt is set when the alloc is in a terminal state
func TestAllocRunner_FinishedAtSet(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, ar := testAllocRunner(t, false)
	ar.allocClientStatus = structs.AllocClientStatusFailed
	alloc := ar.Alloc()
	taskFinishedAt := make(map[string]time.Time)
	require.NotEmpty(alloc.TaskStates)
	for name, s := range alloc.TaskStates {
		require.False(s.FinishedAt.IsZero())
		taskFinishedAt[name] = s.FinishedAt
	}

	// Verify that calling again should not mutate finishedAt
	alloc2 := ar.Alloc()
	for name, s := range alloc2.TaskStates {
		require.Equal(taskFinishedAt[name], s.FinishedAt)
	}

}

// Test that FinisheAt is set when the alloc is in a terminal state
func TestAllocRunner_FinishedAtSet_TaskEvents(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, ar := testAllocRunner(t, false)
	ar.taskStates[ar.alloc.Job.TaskGroups[0].Tasks[0].Name] = &structs.TaskState{State: structs.TaskStateDead, Failed: true}

	alloc := ar.Alloc()
	taskFinishedAt := make(map[string]time.Time)
	require.NotEmpty(alloc.TaskStates)
	for name, s := range alloc.TaskStates {
{'fraction_non_alphanumeric': 0.13947521865889212, 'fraction_numerical': 0.003731778425655977, 'mean_word_length': 4.8221317040054315, 'pattern_counts': {'":': 0, '<': 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': '7361710', 'n_tokens_mistral': 3180, 'n_tokens_neox': 3124, 'n_words': 631}