cmd env: adding an env package to aid in testing.

This change adds an Env object that can be injected into the various commands.
The Env can be configured with an in memory store for testing purposes.

This changes also adds a function to inject arbitrary errors in the in memory store
and fixes the getters to avoid a panic on a nil result.

Author: sravotto

internal/cmd/env/env.go

@@ -0,0 +1,118 @@
+// Copyright 2022 Cockroach Labs Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package env provides the environment for the commands to run.
+package env
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"io"
+	"os"
+
+	"github.com/cockroachlabs/visus/internal/database"
+	"github.com/cockroachlabs/visus/internal/store"
+	"github.com/spf13/cobra"
+)
+
+// Env defines the environment to run the commands.
+type Env struct {
+	testing bool
+	store   store.Store
+	alt     io.Reader
+}
+
+// Default returns the default environment, for real world use.
+func Default() *Env {
+	return &Env{
+		alt: os.Stdin,
+	}
+}
+
+// Testing sets up an environment for testing commands.
+func Testing(ctx context.Context) *Env {
+	store := &store.Memory{}
+	err := store.Init(ctx)
+	if err != nil {
+		// this shouldn't happen for in memory store.
+		panic("failed to initialize the store")
+	}
+	return &Env{
+		testing: true,
+		store:   store,
+	}
+}
+
+// InjectReader sets the alternative reader to use in the env.
+func (e *Env) InjectReader(r io.Reader) {
+	if e.testing {
+		e.alt = r
+	}
+}
+
+// InjectStoreError forces the in memory store to return an error.
+func (e *Env) InjectStoreError(err error) {
+	if mem, ok := e.store.(*store.Memory); e.testing && ok {
+		mem.InjectError(err)
+	}
+}
+
+// ProvideStore returns the Store to use.
+func (e *Env) ProvideStore(ctx context.Context, databaseURL string) (store.Store, error) {
+	if e.testing {
+		return e.store, nil
+	}
+	conn, err := database.New(ctx, databaseURL)
+	if err != nil {
+		return nil, err
+	}
+	e.store = store.New(conn)
+	return e.store, nil
+}
+
+// ReadFile reads the named file. If the filename is "-", it will read
+// the content from the alternate reader in the environment.
+func (e *Env) ReadFile(file string) ([]byte, error) {
+	var data []byte
+	if file == "-" {
+		var buffer bytes.Buffer
+		scanner := bufio.NewScanner(e.alt)
+		for scanner.Scan() {
+			buffer.Write(scanner.Bytes())
+			buffer.WriteString("\n")
+		}
+		if err := scanner.Err(); err != nil {
+			return nil, err
+		}
+		data = buffer.Bytes()
+		return data, nil
+	}
+	return os.ReadFile(file)
+}
+
+// TestCommand executes the given command with the supplied arguments
+// and returns the standard output as a string.
+func (e *Env) TestCommand(f func(*Env) *cobra.Command, args []string) (string, error) {
+	cmd := f(e)
+	buffer := bytes.NewBufferString("")
+	cmd.SetOut(buffer)
+	cmd.SetArgs(args)
+	err := cmd.Execute()
+	if err != nil {
+		return "", err
+	}
+	out, err := io.ReadAll(buffer)
+	return string(out), err
+}

internal/store/memory.go

@@ -16,79 +16,125 @@ package store
 
 import (
 	"context"
+	"slices"
 	"sync"
 	"time"
 )
 
 // Memory stores the configuration in memory. Used for testing.
 type Memory struct {
-	MainNode    bool
 	collections *sync.Map
 	histograms  *sync.Map
 	scans       *sync.Map
+
+	mu struct {
+		sync.RWMutex
+		err      error // Error to return
+		mainNode bool
+	}
 }
 
 var _ Store = &Memory{}
 
 // DeleteCollection implements store.Store.
 func (m *Memory) DeleteCollection(_ context.Context, name string) error {
+	if m.Error() != nil {
+		return m.Error()
+	}
 	m.collections.Delete(name)
 	return nil
 }
 
 // DeleteHistogram implements store.Store.
 func (m *Memory) DeleteHistogram(_ context.Context, name string) error {
+	if m.Error() != nil {
+		return m.Error()
+	}
 	m.histograms.Delete(name)
 	return nil
 }
 
 // DeleteScan implements store.Store.
 func (m *Memory) DeleteScan(_ context.Context, name string) error {
+	if m.Error() != nil {
+		return m.Error()
+	}
 	m.scans.Delete(name)
 	return nil
 }
 
+// Error returns the injected error
+func (m *Memory) Error() error {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.mu.err
+}
+
 // GetCollection implements store.Store.
 func (m *Memory) GetCollection(_ context.Context, name string) (*Collection, error) {
-	res, _ := m.collections.Load(name)
+	if m.Error() != nil {
+		return nil, m.Error()
+	}
+	res, ok := m.collections.Load(name)
+	if !ok {
+		return nil, nil
+	}
 	return res.(*Collection), nil
 }
 
 // GetCollectionNames implements store.Store.
 func (m *Memory) GetCollectionNames(_ context.Context) ([]string, error) {
-	return getNames(m.collections)
+	return m.getNames(m.collections)
 }
 
 // GetHistogram implements store.Store.
 func (m *Memory) GetHistogram(_ context.Context, name string) (*Histogram, error) {
-	res, _ := m.histograms.Load(name)
-	return res.(*Histogram), nil
+	if m.Error() != nil {
+		return nil, m.Error()
+	}
+	res, ok := m.histograms.Load(name)
+	if !ok {
+		return nil, m.Error()
+	}
+	return res.(*Histogram), m.Error()
 }
 
 // GetHistogramNames implements store.Store.
 func (m *Memory) GetHistogramNames(_ context.Context) ([]string, error) {
-	return getNames(m.histograms)
+	return m.getNames(m.histograms)
 }
 
 // GetMetrics implements store.Store.
 func (m *Memory) GetMetrics(ctx context.Context, name string) ([]Metric, error) {
+	if m.Error() != nil {
+		return nil, m.Error()
+	}
 	coll, _ := m.GetCollection(ctx, name)
 	return coll.Metrics, nil
 }
 
 // GetScan implements store.Store.
 func (m *Memory) GetScan(_ context.Context, name string) (*Scan, error) {
-	res, _ := m.scans.Load(name)
+	if m.Error() != nil {
+		return nil, m.Error()
+	}
+	res, ok := m.scans.Load(name)
+	if !ok {
+		return nil, nil
+	}
 	return res.(*Scan), nil
 }
 
 // GetScanNames implements store.Store.
 func (m *Memory) GetScanNames(_ context.Context) ([]string, error) {
-	return getNames(m.scans)
+	return m.getNames(m.scans)
 }
 
 // GetScanPatterns implements store.Store.
 func (m *Memory) GetScanPatterns(ctx context.Context, name string) ([]Pattern, error) {
+	if m.Error() != nil {
+		return nil, m.Error()
+	}
 	scan, _ := m.GetScan(ctx, name)
 	return scan.Patterns, nil
 }
@@ -98,37 +144,67 @@ func (m *Memory) Init(_ context.Context) error {
 	m.collections = &sync.Map{}
 	m.histograms = &sync.Map{}
 	m.scans = &sync.Map{}
+	m.InjectError(nil)
 	return nil
 }
 
+// InjectError sets the error that will be returned on each subsequent call.
+func (m *Memory) InjectError(err error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.mu.err = err
+}
+
 // IsMainNode implements store.Store.
 func (m *Memory) IsMainNode(_ context.Context, lastUpdated time.Time) (bool, error) {
-	return m.MainNode, nil
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.mu.mainNode, m.mu.err
 }
 
 // PutCollection implements store.Store.
 func (m *Memory) PutCollection(_ context.Context, collection *Collection) error {
+	if m.Error() != nil {
+		return m.Error()
+	}
 	m.collections.Store(collection.Name, collection)
 	return nil
 }
 
 // PutHistogram implements store.Store.
 func (m *Memory) PutHistogram(_ context.Context, histogram *Histogram) error {
+	if m.Error() != nil {
+		return m.Error()
+	}
 	m.histograms.Store(histogram.Name, histogram)
 	return nil
 }
 
 // PutScan implements store.Store.
 func (m *Memory) PutScan(_ context.Context, scan *Scan) error {
+	if m.Error() != nil {
+		return m.Error()
+	}
 	m.scans.Store(scan.Name, scan)
 	return nil
 }
 
-func getNames(m *sync.Map) ([]string, error) {
+// SetMainNode sets this store as the main node.
+func (m *Memory) SetMainNode(main bool) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.mu.mainNode = main
+}
+
+func (m *Memory) getNames(vals *sync.Map) ([]string, error) {
+	if m.Error() != nil {
+		return nil, m.Error()
+	}
 	names := make([]string, 0)
-	m.Range(func(key any, value any) bool {
+	vals.Range(func(key any, value any) bool {
 		names = append(names, key.(string))
 		return true
 	})
+	slices.Sort(names)
 	return names, nil
 }