Skip to content
This repository has been archived by the owner on Apr 29, 2020. It is now read-only.

Commit

Permalink
added tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Zongkun Yang committed Jul 3, 2018
1 parent e481c95 commit 0f13199
Show file tree
Hide file tree
Showing 5 changed files with 176 additions and 62 deletions.
2 changes: 1 addition & 1 deletion bin/p2-bootstrap/bootstrap.go
View file
Expand Up @@ -44,7 +44,7 @@ func main() {
nodeName := types.NodeName(hostname)
agentManifest, err := manifest.FromPath(*agentManifestPath)
if err != nil {
log.Fatalln("Could not get agent manifest: %s", err)
log.Fatalf("Could not get agent manifest: %s", err)
}
log.Println("Installing and launching consul")

Expand Down
40 changes: 3 additions & 37 deletions bin/p2-preparer/main.go
View file
Expand Up @@ -5,16 +5,12 @@ import (
"os/signal"
"sync"
"syscall"
"time"

"github.com/Sirupsen/logrus"
"gopkg.in/alecthomas/kingpin.v2"

"github.com/square/p2/pkg/constants"
"github.com/square/p2/pkg/logging"
"github.com/square/p2/pkg/preparer"
"github.com/square/p2/pkg/types"
"github.com/square/p2/pkg/util"
"github.com/square/p2/pkg/util/param"
"github.com/square/p2/pkg/version"
"github.com/square/p2/pkg/watch"
Expand Down Expand Up @@ -89,39 +85,9 @@ func main() {
logger.WithError(err).Fatalf("Could not do initial build reality at launch: %s", err)
}

whiteListPods := make(map[types.PodID]bool)
if preparerConfig.PodWhitelistFile != "" {
// Keep loopinng the white list of pods in pod whitelist file, and install the non existing pods
// exit while the white list is empty
for {
_, err := os.Stat(preparerConfig.PodWhitelistFile)
// if the whitelist file does not exist, end the loop and proceed
if os.IsNotExist(err) {
logger.WithError(err).Warningf("Pod whilelist file does not exist")
break
}

podWhitelist, err := util.LoadTokens(preparerConfig.PodWhitelistFile)
if err != nil {
logger.WithError(err).WithField("path", preparerConfig.PodWhitelistFile).Fatalln("Could not read pod whitelist file")
}

// if the pod white list is empty, jump out of the loop, then p2-preparer proceeds to start with full functionality
if len(podWhitelist) == 0 {
break
}

for pod := range podWhitelist {
whiteListPods[types.PodID(pod)] = true
}

err = prep.InstallWhiteListPods(whiteListPods, preparerConfig)
if err != nil {
logger.WithError(err).Fatalln("Could not install whitelist pods")
}

time.Sleep(constants.P2WhitelistCheckInterval)
}
podWhiteList, err := prep.CheckPodWhitelist(preparerConfig)
if err != nil {
logger.WithError(err).Fatalf("Error occurs when checking pod whitelist %+v", podWhiteList)
}

go prep.WatchForPodManifestsForNode(quitMainUpdate)
Expand Down
2 changes: 1 addition & 1 deletion integration/single-node-slug-deploy/check.go
View file
Expand Up @@ -101,7 +101,7 @@ func main() {

requireFile, err := createRequireFile(tempdir)
if err != nil {
log.Fatalln("Could not create temp require file, bailing: %s", err)
log.Fatalf("Could not create temp require file, bailing: %s", err)
}

userHookManifest, err := userCreationHookManifest(tempdir, noAddUser)
Expand Down
86 changes: 65 additions & 21 deletions pkg/preparer/orchestrate.go
View file
Expand Up @@ -5,6 +5,8 @@ import (
"fmt"
"time"

"os"

"github.com/Sirupsen/logrus"
"github.com/square/p2/pkg/artifact"
"github.com/square/p2/pkg/auth"
Expand Down Expand Up @@ -79,16 +81,57 @@ func (p podWorkerID) String() string {
return fmt.Sprintf("%s-%s", p.podID.String(), p.podUniqueKey)
}

func (p *Preparer) InstallWhiteListPods(whiteListPods map[types.PodID]bool, preparerConfig *PreparerConfig) error {
func (p *Preparer) CheckPodWhitelist(preparerConfig *PreparerConfig) (map[types.PodID]bool, error) {
whiteListPods := make(map[types.PodID]bool)
if preparerConfig.PodWhitelistFile != "" {
// Keep loopinng the white list of pods in pod whitelist file, and install the non existing pods
// exit while the white list is empty
for {
_, err := os.Stat(preparerConfig.PodWhitelistFile)
// if the whitelist file does not exist, end the loop and proceed
if os.IsNotExist(err) {
p.Logger.WithError(err).Warningf("Pod whilelist file does not exist")
break
}

podWhitelist, err := util.LoadTokens(preparerConfig.PodWhitelistFile)
if err != nil {
return nil, err
}

// if the pod white list is empty, jump out of the loop, then p2-preparer proceeds to start with full functionality
if len(podWhitelist) == 0 {
break
}

for _, pod := range podWhitelist {
whiteListPods[types.PodID(pod)] = true
}
p.Logger.WithField("path", preparerConfig.PodWhitelistFile).Printf("Whitelist pods: %+v", whiteListPods)
err = p.installWhiteListPods(whiteListPods)
if err != nil {
return whiteListPods, err
}
p.Logger.Println("All pods in whitelist, have been successfully installed, empty or remove the pod whitelist file to proceed.")
time.Sleep(constants.P2WhitelistCheckInterval)
}
}
return whiteListPods, nil
}

func (p *Preparer) installWhiteListPods(whiteListPods map[types.PodID]bool) error {
intentResults, _, err := p.store.ListPods(consul.INTENT_TREE, p.node)
realityResults, _, err := p.store.ListPods(consul.REALITY_TREE, p.node)
var pairs []*ManifestPair
if err != nil {
p.Logger.WithError(err).Errorln("could not check intent")
}
realityResults, _, err := p.store.ListPods(consul.REALITY_TREE, p.node)
if err != nil {
p.Logger.WithError(err).Errorln("could not check reality")
}

var pairs []*ManifestPair
errorChan := make(chan error)
quit := make(chan struct{})

for _, intentResult := range intentResults {
if whiteListPods[intentResult.Manifest.ID()] {
manifestPair := &ManifestPair{
Expand All @@ -106,11 +149,11 @@ func (p *Preparer) InstallWhiteListPods(whiteListPods map[types.PodID]bool, prep
pairs = append(pairs, manifestPair)
}
}

p.Logger.WithField("whitelistPodToInstall", pairs).Println("Pods to be installed")
for _, pair := range pairs {
go p.handleWhiteListPod(pair, errorChan, quit)
}
for i := 0; i < len(pairs); i++ {
for range pairs {
select {
case <-quit:
case err := <-errorChan:
Expand All @@ -120,39 +163,40 @@ func (p *Preparer) InstallWhiteListPods(whiteListPods map[types.PodID]bool, prep
return nil
}

func (p *Preparer) handleWhiteListPod(intentManifest *ManifestPair, errorChan chan<- error, quit chan<- struct{}) {
func (p *Preparer) handleWhiteListPod(pair *ManifestPair, errorChan chan<- error, quit chan<- struct{}) {
var pod *pods.Pod
var err error
var manifestLogger logging.Logger
sha, _ := intentManifest.Intent.SHA()
sha, _ := pair.Intent.SHA()

manifestLogger = p.Logger.SubLogger(logrus.Fields{
"pod": intentManifest.ID,
"pod": pair.ID,
"sha": sha,
"pod_unique_key": intentManifest.PodUniqueKey,
"pod_unique_key": pair.PodUniqueKey,
})
manifestLogger.NoFields().Debugln("processing whitelist pod manifest")

if intentManifest.PodUniqueKey == "" {
pod = p.podFactory.NewLegacyPod(intentManifest.ID)
if pair.PodUniqueKey == "" {
pod = p.podFactory.NewLegacyPod(pair.ID)
} else {
pod, err = p.podFactory.NewUUIDPod(intentManifest.ID, intentManifest.PodUniqueKey)
pod, err = p.podFactory.NewUUIDPod(pair.ID, pair.PodUniqueKey)
if err != nil {
manifestLogger.WithError(err).Errorln("Could not initialize pod")
errorChan <- util.Errorf("failed to initialize pod: %s, error: %v", intentManifest.Intent.ID(), err)
errorChan <- util.Errorf("failed to initialize pod: %s, error: %v", pair.Intent.ID(), err)
}
}

err = p.preparePod(*intentManifest, pod, manifestLogger)
p.Logger.WithField("podManifest", pair).Println("Start installing whitelist pod")
err = p.preparePod(pair, pod, manifestLogger)
if err != nil {
errorChan <- util.Errorf("failed to install pod: %s, error: %v", intentManifest.Intent.ID(), err)
errorChan <- util.Errorf("failed to install pod: %s, error: %v", pair.Intent.ID(), err)
return
}
ok := p.resolvePair(*intentManifest, pod, manifestLogger)
ok := p.resolvePair(*pair, pod, manifestLogger)
if !ok {
errorChan <- util.Errorf("failed to install pod: %s", intentManifest.Intent.ID())
errorChan <- util.Errorf("failed to install pod: %s", pair.Intent.ID())
return
}
p.Logger.WithField("podManifest", pair).Println("Finished installation of the whitelist pod")
quit <- struct{}{}
}

Expand Down Expand Up @@ -298,7 +342,7 @@ func (p *Preparer) handlePods(podChan <-chan ManifestPair, quit <-chan struct{})
break
}
}
err = p.preparePod(nextLaunch, pod, manifestLogger)
err = p.preparePod(&nextLaunch, pod, manifestLogger)
if err != nil {
break
}
Expand All @@ -320,7 +364,7 @@ func (p *Preparer) handlePods(podChan <-chan ManifestPair, quit <-chan struct{})
}
}

func (p *Preparer) preparePod(nextLaunch ManifestPair, pod *pods.Pod, manifestLogger logging.Logger) error {
func (p *Preparer) preparePod(nextLaunch *ManifestPair, pod *pods.Pod, manifestLogger logging.Logger) error {
// TODO better solution: force the preparer to have a 0s default timeout, prevent KILLs
if pod.Id == constants.PreparerPodID {
pod.DefaultTimeout = time.Duration(0)
Expand Down
108 changes: 106 additions & 2 deletions pkg/preparer/orchestrate_test.go
View file
Expand Up @@ -9,20 +9,28 @@ import (
"testing"
"time"

"strings"

"path/filepath"

. "github.com/anthonybishopric/gotcha"
"github.com/hashicorp/consul/api"
"github.com/square/p2/pkg/artifact"
"github.com/square/p2/pkg/auth"
"github.com/square/p2/pkg/constants"
"github.com/square/p2/pkg/hooks"
"github.com/square/p2/pkg/logging"
"github.com/square/p2/pkg/manifest"
"github.com/square/p2/pkg/store/consul"
"github.com/square/p2/pkg/store/consul/statusstore"
"github.com/square/p2/pkg/store/consul/statusstore/podstatus"
"github.com/square/p2/pkg/types"
"github.com/square/p2/pkg/util"
"github.com/square/p2/pkg/util/size"
"golang.org/x/crypto/openpgp"
"golang.org/x/crypto/openpgp/clearsign"
"golang.org/x/crypto/openpgp/packet"
context "golang.org/x/net/context"
)

type TestPod struct {
Expand Down Expand Up @@ -174,17 +182,37 @@ func testSignedManifest(t *testing.T, modify func(manifest.Builder, *openpgp.Ent
type FakeStore struct {
currentManifest manifest.Manifest
currentManifestError error
wipedReality bool
}

type FakeStatusStore struct {
}

func (f *FakeStatusStore) Get(types.PodUniqueKey) (podstatus.PodStatus, *api.QueryMeta, error) {
return podstatus.PodStatus{}, nil, statusstore.NoStatusError{}
}

func (f *FakeStore) ListPods(consul.PodPrefix, types.NodeName) ([]consul.ManifestResult, time.Duration, error) {
func (f *FakeStatusStore) MutateStatus(context.Context, types.PodUniqueKey, func(podstatus.PodStatus) (podstatus.PodStatus, error)) error {
return nil
}

func (f *FakeStore) ListPods(prefix consul.PodPrefix, node types.NodeName) ([]consul.ManifestResult, time.Duration, error) {
if f.currentManifest == nil {
return nil, 0, nil
}
if f.currentManifestError != nil {
return nil, 0, f.currentManifestError
}
if f.wipedReality && prefix == consul.REALITY_TREE {
return nil, 0, nil
}
uuid := types.PodUniqueKey("1")
manifest := &consul.ManifestResult{
Manifest: f.currentManifest,
PodUniqueKey: uuid,
}
return []consul.ManifestResult{
{Manifest: f.currentManifest},
*manifest,
}, 0, nil
}

Expand Down Expand Up @@ -218,9 +246,85 @@ func testPreparer(t *testing.T, f *FakeStore) (*Preparer, *fakeHooks, string) {
hooks := &fakeHooks{}
p.hooks = hooks
p.store = f
p.podStatusStore = &FakeStatusStore{}
return p, hooks, podRoot
}

func TestPrepareSkipInstallIfPodWhiteListFileNotExist(t *testing.T) {
p, _, _ := testPreparer(t, &FakeStore{})
_, filename, _, _ := runtime.Caller(0)
testPath := filepath.Join(filepath.Dir(filename), "pod_whitelist_dummy")
cfg := &PreparerConfig{
PodWhitelistFile: testPath,
}
whitelistPods, err := p.CheckPodWhitelist(cfg)
Assert(t).IsTrue(len(whitelistPods) == 0, "With no existence of pod whitelist file, the list should be empty")
Assert(t).IsNil(err, "Should not have erred when pod whitelist file does not exists")
}

func TestPrepareCheckWhiteListFileIfConfigured(t *testing.T) {
p, _, _ := testPreparer(t, &FakeStore{})
_, filename, _, _ := runtime.Caller(0)
testPath := filepath.Join(filepath.Dir(filename), "pod_whitelist")

file, _ := os.Create(testPath)
defer file.Close()

file.Write([]byte("slug"))

cfg := &PreparerConfig{
PodWhitelistFile: testPath,
}
defer os.Remove(testPath)
go func(cfg *PreparerConfig) {
whitelistPods, err := p.CheckPodWhitelist(cfg)
Assert(t).IsTrue(len(whitelistPods) == 1, "Successfully read the pods from whitelist file")
Assert(t).IsNil(err, "Successfully exited the infinite for loop of checking pod whitelist, when the file is deleted")
}(cfg)

time.Sleep(time.Second)
}

func TestPrepareInstallWhiteListPodsInTheIntentSkipInstall(t *testing.T) {
whiteListPods := make(map[types.PodID]bool)
whiteListPods["hello"] = true

store := &FakeStore{
currentManifest: testManifest(t),
wipedReality: true,
}

p, _, fakePodRoot := testPreparer(t, store)
defer p.Close()
defer os.RemoveAll(fakePodRoot)

// error out the podID to be installed here, the further installation is not included here, tested by other tests already.
// We are verifying the podID to be installed
err := p.installWhiteListPods(whiteListPods)
s := strings.Split(err.Error(), " ")
podID := s[len(s)-1]
Assert(t).AreEqual(podID, "hello", "The podID to install should be the one in intent")
}

func TestPrepareWillNotInstallWhiteListPodsNotInTheIntentSkipInstall(t *testing.T) {
whiteListPods := make(map[types.PodID]bool)
whiteListPods["test"] = true

store := &FakeStore{
currentManifest: testManifest(t),
wipedReality: true,
}

p, _, fakePodRoot := testPreparer(t, store)
defer p.Close()
defer os.RemoveAll(fakePodRoot)

// error out the podID to be installed here, we are stipping installation by returning an err, if err is not nil,
// means there is no pod to install, which is expected
err := p.installWhiteListPods(whiteListPods)
Assert(t).IsTrue(err == nil, "InstallAndLaunchPod not executed")
}

func TestPreparerLaunchesNewPodsThatArentInstalledYet(t *testing.T) {
testPod := &TestPod{
launchSuccess: true,
Expand Down

0 comments on commit 0f13199

Please sign in to comment.