DDI/agent/core/cli_test.go

// SPDX-License-Identifier: Apache-2.0

package core

import (
	"bytes"
	"context"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func TestMainBasics(t *testing.T) {
	tests := []struct {
		name       string
		args       []string
		wantCode   int
		wantStdout string
		wantStderr string
	}{
		{name: "no args", args: nil, wantCode: 2, wantStderr: "Usage: ddi-agent-test"},
		{name: "version", args: []string{"version"}, wantCode: 0, wantStdout: "ddi-agent-test "},
		{name: "--version", args: []string{"--version"}, wantCode: 0, wantStdout: "ddi-agent-test "},
		{name: "help", args: []string{"help"}, wantCode: 0, wantStdout: "Agent role: dns"},
		{name: "help lists enroll", args: []string{"--help"}, wantCode: 0, wantStdout: "--ca-hash sha256:<hex>"},
		{name: "unknown", args: []string{"explode"}, wantCode: 2, wantStderr: `unknown command "explode"`},
		{name: "enroll without flags", args: []string{"enroll"}, wantCode: 2, wantStderr: "enroll needs --server"},
		{name: "bad flag", args: []string{"run", "--bogus"}, wantCode: 2, wantStderr: "flag provided but not defined"},
		{name: "extra argument", args: []string{"status", "extra"}, wantCode: 2, wantStderr: `unexpected argument "extra"`},
		{name: "flag help", args: []string{"status", "-h"}, wantCode: 0, wantStderr: "-config"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var out, errOut bytes.Buffer
			code := Main("ddi-agent-test", KindDNS, nil, tt.args, &out, &errOut)
			if code != tt.wantCode {
				t.Errorf("code = %d, want %d (stderr %q)", code, tt.wantCode, errOut.String())
			}
			if tt.wantStdout != "" && !strings.Contains(out.String(), tt.wantStdout) {
				t.Errorf("stdout %q does not contain %q", out.String(), tt.wantStdout)
			}
			if tt.wantStderr != "" && !strings.Contains(errOut.String(), tt.wantStderr) {
				t.Errorf("stderr %q does not contain %q", errOut.String(), tt.wantStderr)
			}
		})
	}
}

// writeConfig writes a configuration with its own state directory.
func writeConfig(t *testing.T, extra string) (cfgPath, stateDir string) {
	t.Helper()
	dir := t.TempDir()
	stateDir = filepath.Join(dir, "state")
	cfgPath = filepath.Join(dir, "agent.yaml")
	data := "state_dir: " + filepath.ToSlash(stateDir) + "\nlog:\n  format: text\n" + extra
	if err := os.WriteFile(cfgPath, []byte(data), 0o600); err != nil {
		t.Fatal(err)
	}
	return cfgPath, stateDir
}

func TestMainEnrollStatusRun(t *testing.T) {
	srv := newFakeServer(t, KindDNS)
	addr := srv.start(t)
	cfgPath, stateDir := writeConfig(t, "dns:\n  zone_dir: /var/lib/bind/x\n")

	var out, errOut bytes.Buffer
	if code := Main("ddi-agent-dns", KindDNS, nil, []string{"status", "--config", cfgPath}, &out, &errOut); code != 0 {
		t.Fatalf("status before enroll = %d: %s", code, errOut.String())
	}
	if !strings.Contains(out.String(), "not enrolled") {
		t.Errorf("status = %q", out.String())
	}

	// run without identity fails.
	out.Reset()
	errOut.Reset()
	factory := func(decode func(any) error) (Role, error) { return newFakeRole(KindDNS), nil }
	if code := Main("ddi-agent-dns", KindDNS, factory, []string{"run", "--config", cfgPath}, &out, &errOut); code != 1 {
		t.Fatalf("run before enroll = %d", code)
	}
	if !strings.Contains(errOut.String(), "not enrolled") {
		t.Errorf("stderr = %q", errOut.String())
	}

	args := []string{"enroll", "--config", cfgPath, "--server", addr, "--token", srv.token,
		"--ca-hash", srv.fingerprint(), "--server-name", "localhost"}
	out.Reset()
	errOut.Reset()
	if code := Main("ddi-agent-dns", KindDNS, nil, args, &out, &errOut); code != 0 {
		t.Fatalf("enroll = %d: %s", code, errOut.String())
	}
	if !strings.Contains(out.String(), srv.agentID) {
		t.Errorf("enroll output = %q", out.String())
	}
	out.Reset()
	errOut.Reset()
	if code := Main("ddi-agent-dns", KindDNS, nil, args, &out, &errOut); code != 1 || !strings.Contains(errOut.String(), "--force") {
		t.Errorf("second enroll = %d: %s", code, errOut.String())
	}

	// Run the agent until revision 1 is reported.
	srv.with(func(f *fakeServer) { f.desired = 1 })
	cfg, err := LoadConfig(cfgPath, KindDNS)
	if err != nil {
		t.Fatal(err)
	}
	var gotSection struct {
		ZoneDir string `yaml:"zone_dir"`
	}
	role := newFakeRole(KindDNS)
	factory = func(decode func(any) error) (Role, error) {
		if err := decode(&gotSection); err != nil {
			return nil, err
		}
		return role, nil
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- runAgent(ctx, KindDNS, factory, cfg, testLogger()) }()
	waitFor(t, "deployment report", func() bool { return len(serverReports(srv)) == 1 })
	cancel()
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("runAgent = %v", err)
		}
	case <-time.After(10 * time.Second):
		t.Fatal("runAgent did not return")
	}
	if gotSection.ZoneDir != "/var/lib/bind/x" {
		t.Errorf("role section = %+v", gotSection)
	}

	out.Reset()
	errOut.Reset()
	if code := Main("ddi-agent-dns", KindDNS, nil, []string{"status", "--config", cfgPath}, &out, &errOut); code != 0 {
		t.Fatalf("status = %d: %s", code, errOut.String())
	}
	for _, want := range []string{srv.agentID, "Applied revision:", "revision 1 succeeded", "Stored revisions:    1", filepath.ToSlash(stateDir)} {
		if !strings.Contains(out.String(), want) {
			t.Errorf("status lacks %q:\n%s", want, out.String())
		}
	}
}

func TestRunAgentErrors(t *testing.T) {
	cfgPath, _ := writeConfig(t, "")
	cfg, err := LoadConfig(cfgPath, KindDNS)
	if err != nil {
		t.Fatal(err)
	}
	if err := runAgent(context.Background(), KindDNS, nil, cfg, testLogger()); !errors.Is(err, errRoleNotAvailable) {
		t.Errorf("nil factory: %v", err)
	}
	// Enrolled as a DHCP agent, run as DNS agent.
	store, err := OpenStore(cfg.StateDir)
	if err != nil {
		t.Fatal(err)
	}
	pki := newTestPKI(t)
	id, key := testIdentity(t, pki)
	id.Role = KindDHCP
	if err := store.SaveIdentity(id, key, false); err != nil {
		t.Fatal(err)
	}
	factory := func(func(any) error) (Role, error) { return newFakeRole(KindDNS), nil }
	if err := runAgent(context.Background(), KindDNS, factory, cfg, testLogger()); err == nil || !strings.Contains(err.Error(), "dhcp agent") {
		t.Errorf("role mismatch: %v", err)
	}
	if err := runAgent(context.Background(), KindDHCP, func(func(any) error) (Role, error) {
		return nil, errors.New("kea missing")
	}, cfg, testLogger()); err == nil || !strings.Contains(err.Error(), "kea missing") {
		t.Errorf("factory error: %v", err)
	}
}

func TestMainConfigErrors(t *testing.T) {
	cfgPath, _ := writeConfig(t, "bogus: 1\n")
	for _, args := range [][]string{
		{"run", "--config", cfgPath},
		{"status", "--config", cfgPath},
		{"enroll", "--config", cfgPath, "--server", "x:1", "--token", "a.b", "--ca-hash", "sha256:00"},
	} {
		var out, errOut bytes.Buffer
		if code := Main("ddi-agent-dns", KindDNS, nil, args, &out, &errOut); code != 1 || !strings.Contains(errOut.String(), "bogus") {
			t.Errorf("%v: code %d, stderr %q", args[0], code, errOut.String())
		}
	}
}