theta-agent/config_test.go

package main

import (
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestLoadConfig(t *testing.T) {
	// Create a temporary directory for config files
	tmpDir, err := os.MkdirTemp("", "agent-config-test")
	if err != nil {
		t.Fatalf("failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	tests := []struct {
		name        string
		yamlContent string
		filename    string
		expectErr   bool
	}{
		{
			name: "valid config",
			yamlContent: `
server_url: "http://sso.local"
auth_token: "secret-token"
location: "datacenter-1"
capabilities:
  telemetry: true
  configure_ldap: true
  reboot: false
  service_control: ["nginx", "gitea"]
  arbitrary_bash: false
`,
			filename:  "valid.yml",
			expectErr: false,
		},
		{
			name:        "invalid yaml",
			yamlContent: "invalid: [yaml: content",
			filename:    "invalid.yml",
			expectErr:   true,
		},
		{
			name:        "missing file",
			yamlContent: "",
			filename:    "nonexistent.yml",
			expectErr:   true,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			path := filepath.Join(tmpDir, tc.filename)
			if tc.yamlContent != "" {
				err := os.WriteFile(path, []byte(tc.yamlContent), 0644)
				if err != nil {
					t.Fatalf("failed to write temp file: %v", err)
				}
			}

			cfg, err := LoadConfig(path)
			if (err != nil) != tc.expectErr {
				t.Errorf("LoadConfig() error = %v, expectErr %v", err, tc.expectErr)
				return
			}

			if !tc.expectErr && cfg == nil {
				t.Error("LoadConfig() returned nil config without error")
			}
		})
	}
}

func TestCanManageService(t *testing.T) {
	caps := Capabilities{
		ServiceControl: []string{"nginx", "gitea"},
	}

	tests := []struct {
		service  string
		expected bool
	}{
		{"nginx", true},
		{"gitea", true},
		{"ssh", false},
		{"", false},
	}

	for _, tc := range tests {
		t.Run(tc.service, func(t *testing.T) {
			if got := caps.CanManageService(tc.service); got != tc.expected {
				t.Errorf("CanManageService(%q) = %v, want %v", tc.service, got, tc.expected)
			}
		})
	}
}

// PersistEnrollment rewrites a hand-edited file, so it must replace exactly the
// credential lines and leave everything else -- comments, capabilities,
// formatting -- untouched.
func TestPersistEnrollmentPreservesFile(t *testing.T) {
	dir := t.TempDir()
	path := dir + "/agent.yml"
	original := `# theta-agent configuration file
server_url: "https://sso.example.com"

# Bootstrap credential, exchanged on first connect.
join_key: "tjk_abc123"
public_key: ""
location: "rack-4"

capabilities:
  telemetry: true
  # keep this comment
  reboot: false
  service_control: ["nginx"]
`
	if err := os.WriteFile(path, []byte(original), 0600); err != nil {
		t.Fatal(err)
	}
	cm, err := NewConfigManager(path)
	if err != nil {
		t.Fatal(err)
	}

	if err := cm.PersistEnrollment("issued-token-xyz", "PUBKEYBASE64"); err != nil {
		t.Fatalf("PersistEnrollment: %v", err)
	}

	out, _ := os.ReadFile(path)
	got := string(out)

	for _, want := range []string{
		`auth_token: "issued-token-xyz"`,
		`public_key: "PUBKEYBASE64"`,
		`join_key: ""`, // blanked: a fleet-wide key must not linger once unneeded
		"# theta-agent configuration file",
		"# keep this comment",
		`location: "rack-4"`,
		`service_control: ["nginx"]`,
	} {
		if !strings.Contains(got, want) {
			t.Errorf("expected %q in rewritten config, got:\n%s", want, got)
		}
	}

	// and the in-memory config is live without a restart
	if cm.Get().AuthToken != "issued-token-xyz" {
		t.Errorf("config not reloaded: AuthToken = %q", cm.Get().AuthToken)
	}
	if cm.Get().Credential() != "issued-token-xyz" {
		t.Errorf("Credential() should prefer the issued token, got %q", cm.Get().Credential())
	}

	// file must stay 0600 -- it now holds a credential
	fi, _ := os.Stat(path)
	if fi.Mode().Perm() != 0600 {
		t.Errorf("expected mode 0600, got %o", fi.Mode().Perm())
	}
}

func TestPersistEnrollmentAddsMissingKeys(t *testing.T) {
	dir := t.TempDir()
	path := dir + "/agent.yml"
	// No auth_token or public_key lines at all.
	if err := os.WriteFile(path, []byte("server_url: \"https://sso.example.com\"\njoin_key: \"tjk_x\"\n"), 0600); err != nil {
		t.Fatal(err)
	}
	cm, err := NewConfigManager(path)
	if err != nil {
		t.Fatal(err)
	}
	if err := cm.PersistEnrollment("tok", "pk"); err != nil {
		t.Fatalf("PersistEnrollment: %v", err)
	}
	if cm.Get().AuthToken != "tok" || cm.Get().PublicKey != "pk" {
		out, _ := os.ReadFile(path)
		t.Errorf("keys not appended; file:\n%s", out)
	}
}

func TestCredentialPrefersAuthToken(t *testing.T) {
	c := &Config{JoinKey: "tjk_x"}
	if c.Credential() != "tjk_x" {
		t.Errorf("unenrolled agent should present the join key, got %q", c.Credential())
	}
	c.AuthToken = "own-token"
	if c.Credential() != "own-token" {
		t.Errorf("enrolled agent must present its own token, not the join key, got %q", c.Credential())
	}
}

func TestPersistEnrollmentRejectsEmptyToken(t *testing.T) {
	dir := t.TempDir()
	path := dir + "/agent.yml"
	os.WriteFile(path, []byte("server_url: \"x\"\n"), 0600)
	cm, _ := NewConfigManager(path)
	if err := cm.PersistEnrollment("", "pk"); err == nil {
		t.Error("expected an error when the server sends no token")
	}
}