zlh-agent/internal/alloy/alloy_test.go

package alloy

import (
	"errors"
	"net"
	"os"
	"strings"
	"testing"
	"time"

	"zlh-agent/internal/state"
)

const templateConfig = `logging {
  level = "info"
}

prometheus.exporter.unix "local" {}

prometheus.scrape "local" {
  targets    = prometheus.exporter.unix.local.targets
  forward_to = [prometheus.remote_write.zlh_monitor.receiver]
}

prometheus.remote_write "zlh_monitor" {
  endpoint {
    url = "http://10.60.0.25:9090/api/v1/write"
  }

  external_labels = {
    job       = "old",
    instance  = "old",
    collector = "old",
    role      = "old",
    vmid      = "old",
  }
}
`

func TestReplaceExternalLabelsBlockPreservesTemplate(t *testing.T) {
	labels := map[string]string{
		"job":       "integrations/unix",
		"instance":  "10.200.0.46:12345",
		"collector": "alloy",
		"role":      "game-container",
		"vmid":      "5173",
	}

	rendered, err := ReplaceExternalLabelsBlock(templateConfig, labels)
	if err != nil {
		t.Fatalf("ReplaceExternalLabelsBlock: %v", err)
	}

	if !strings.Contains(rendered, `url = "http://10.60.0.25:9090/api/v1/write"`) {
		t.Fatalf("rendered config did not preserve template body:\n%s", rendered)
	}
	if !strings.Contains(rendered, `job       = "integrations/unix",`) {
		t.Fatalf("rendered config missing job label:\n%s", rendered)
	}
	if !strings.Contains(rendered, `instance  = "10.200.0.46:12345",`) {
		t.Fatalf("rendered config missing instance label:\n%s", rendered)
	}
	if strings.Contains(rendered, `job       = "old"`) {
		t.Fatalf("rendered config retained old labels:\n%s", rendered)
	}
}

func TestLabelsDevAndGame(t *testing.T) {
	restoreTestHooks(t)

	devLabels, err := Labels(state.Config{
		VMID:          6001,
		ContainerIP:   "10.60.0.223",
		ContainerType: "dev",
	})
	if err != nil {
		t.Fatalf("Labels dev: %v", err)
	}
	assertLabel(t, devLabels, "instance", "10.60.0.223:12345")
	assertLabel(t, devLabels, "role", "dev-container")
	assertLabel(t, devLabels, "vmid", "6001")

	gameLabels, err := Labels(state.Config{
		VMID:          5001,
		ContainerIP:   "10.60.0.224",
		ContainerType: "game",
	})
	if err != nil {
		t.Fatalf("Labels game: %v", err)
	}
	assertLabel(t, gameLabels, "instance", "10.60.0.224:12345")
	assertLabel(t, gameLabels, "role", "game-container")
	assertLabel(t, gameLabels, "vmid", "5001")
}

func TestEnsureConfigUnchangedDoesNotRestart(t *testing.T) {
	restoreTestHooks(t)

	cfg := state.Config{
		VMID:          6001,
		ContainerIP:   "10.60.0.223",
		ContainerType: "dev",
	}
	labels, err := Labels(cfg)
	if err != nil {
		t.Fatalf("Labels: %v", err)
	}
	rendered, err := ReplaceExternalLabelsBlock(templateConfig, labels)
	if err != nil {
		t.Fatalf("ReplaceExternalLabelsBlock: %v", err)
	}

	restartCalls := 0
	removeCalls := 0
	writeFile = func(path string, data []byte, perm os.FileMode) error {
		if path != tmpConfigPath {
			t.Fatalf("write path = %q, want %q", path, tmpConfigPath)
		}
		return nil
	}
	readFile = func(path string) ([]byte, error) {
		if path != ConfigPath {
			t.Fatalf("read path = %q, want %q", path, ConfigPath)
		}
		return []byte(rendered), nil
	}
	removeFile = func(path string) error {
		removeCalls++
		if path != tmpConfigPath {
			t.Fatalf("remove path = %q, want %q", path, tmpConfigPath)
		}
		return nil
	}
	renameFile = func(string, string) error {
		t.Fatalf("rename should not be called for unchanged config")
		return nil
	}
	runCommand = func(string, ...string) error {
		restartCalls++
		return nil
	}

	result, err := EnsureConfig(cfg)
	if err != nil {
		t.Fatalf("EnsureConfig: %v", err)
	}
	if result.Applied {
		t.Fatalf("Applied = true, want false")
	}
	if removeCalls != 1 {
		t.Fatalf("removeCalls = %d, want 1", removeCalls)
	}
	if restartCalls != 0 {
		t.Fatalf("restartCalls = %d, want 0", restartCalls)
	}
}

func TestEnsureConfigChangedRestartsAndValidates(t *testing.T) {
	restoreTestHooks(t)

	restartCalls := 0
	activeChecks := 0
	listenChecks := 0
	writeFile = func(path string, data []byte, perm os.FileMode) error {
		if path != tmpConfigPath {
			t.Fatalf("write path = %q, want %q", path, tmpConfigPath)
		}
		if !strings.Contains(string(data), `role      = "game-container",`) {
			t.Fatalf("temp config missing rendered labels:\n%s", string(data))
		}
		return nil
	}
	readFile = func(path string) ([]byte, error) {
		return []byte(templateConfig), nil
	}
	removeFile = func(string) error { return nil }
	renameFile = func(oldPath, newPath string) error {
		if oldPath != tmpConfigPath || newPath != ConfigPath {
			t.Fatalf("rename = %q -> %q, want %q -> %q", oldPath, newPath, tmpConfigPath, ConfigPath)
		}
		return nil
	}
	runCommand = func(name string, args ...string) error {
		if name != "systemctl" {
			t.Fatalf("command name = %q, want systemctl", name)
		}
		joined := strings.Join(args, " ")
		switch joined {
		case "restart alloy":
			restartCalls++
		case "is-active --quiet alloy":
			activeChecks++
		default:
			t.Fatalf("unexpected systemctl args: %s", joined)
		}
		return nil
	}
	dialTimeout = func(network, address string, timeout time.Duration) (net.Conn, error) {
		if network != "tcp" || address != "127.0.0.1:12345" {
			t.Fatalf("dial = %s %s, want tcp 127.0.0.1:12345", network, address)
		}
		listenChecks++
		left, right := net.Pipe()
		_ = right.Close()
		return left, nil
	}

	result, err := EnsureConfig(state.Config{
		VMID:          5001,
		ContainerIP:   "10.60.0.224",
		ContainerType: "game",
	})
	if err != nil {
		t.Fatalf("EnsureConfig: %v", err)
	}
	if !result.Applied {
		t.Fatalf("Applied = false, want true")
	}
	if restartCalls != 1 {
		t.Fatalf("restartCalls = %d, want 1", restartCalls)
	}
	if activeChecks != 1 {
		t.Fatalf("activeChecks = %d, want 1", activeChecks)
	}
	if listenChecks != 1 {
		t.Fatalf("listenChecks = %d, want 1", listenChecks)
	}
}

func TestEnsureConfigRetriesFailedUpdate(t *testing.T) {
	restoreTestHooks(t)

	attempts := 0
	readFile = func(path string) ([]byte, error) {
		attempts++
		if attempts == 1 {
			return nil, errors.New("temporary read failure")
		}
		return []byte(templateConfig), nil
	}
	writeFile = func(string, []byte, os.FileMode) error { return nil }
	removeFile = func(string) error { return nil }
	renameFile = func(string, string) error { return nil }
	runCommand = func(string, ...string) error { return nil }

	result, err := EnsureConfig(state.Config{
		VMID:          5001,
		ContainerIP:   "10.60.0.224",
		ContainerType: "game",
	})
	if err != nil {
		t.Fatalf("EnsureConfig: %v", err)
	}
	if !result.Applied {
		t.Fatalf("Applied = false, want true after retry")
	}
	if attempts != 2 {
		t.Fatalf("attempts = %d, want 2", attempts)
	}
}

func TestEnsureConfigRetriesRestartAfterFileChanged(t *testing.T) {
	restoreTestHooks(t)

	restarts := 0
	readFile = func(path string) ([]byte, error) {
		if restarts == 0 {
			return []byte(templateConfig), nil
		}
		labels, err := Labels(state.Config{VMID: 5001, ContainerIP: "10.60.0.224", ContainerType: "game"})
		if err != nil {
			t.Fatalf("Labels: %v", err)
		}
		rendered, err := ReplaceExternalLabelsBlock(templateConfig, labels)
		if err != nil {
			t.Fatalf("ReplaceExternalLabelsBlock: %v", err)
		}
		return []byte(rendered), nil
	}
	writeFile = func(string, []byte, os.FileMode) error { return nil }
	removeFile = func(string) error { return nil }
	renameFile = func(string, string) error { return nil }
	runCommand = func(name string, args ...string) error {
		if name == "systemctl" && strings.Join(args, " ") == "restart alloy" {
			restarts++
			if restarts == 1 {
				return errors.New("temporary restart failure")
			}
		}
		return nil
	}

	result, err := EnsureConfig(state.Config{
		VMID:          5001,
		ContainerIP:   "10.60.0.224",
		ContainerType: "game",
	})
	if err != nil {
		t.Fatalf("EnsureConfig: %v", err)
	}
	if !result.Applied {
		t.Fatalf("Applied = false, want true after restart retry")
	}
	if restarts != 2 {
		t.Fatalf("restarts = %d, want 2", restarts)
	}
}

func assertLabel(t *testing.T, labels map[string]string, key, want string) {
	t.Helper()
	if got := labels[key]; got != want {
		t.Fatalf("labels[%q] = %q, want %q", key, got, want)
	}
}

func restoreTestHooks(t *testing.T) {
	t.Helper()

	oldLocalIP := localIPFunc
	oldRunCommand := runCommand
	oldDialTimeout := dialTimeout
	oldSleep := sleepFunc
	oldWriteFile := writeFile
	oldReadFile := readFile
	oldRemoveFile := removeFile
	oldRenameFile := renameFile

	t.Cleanup(func() {
		localIPFunc = oldLocalIP
		runCommand = oldRunCommand
		dialTimeout = oldDialTimeout
		sleepFunc = oldSleep
		writeFile = oldWriteFile
		readFile = oldReadFile
		removeFile = oldRemoveFile
		renameFile = oldRenameFile
	})

	localIPFunc = func() (string, error) { return "10.60.0.200", nil }
	runCommand = func(string, ...string) error { return nil }
	dialTimeout = func(string, string, time.Duration) (net.Conn, error) {
		left, right := net.Pipe()
		_ = right.Close()
		return left, nil
	}
	sleepFunc = func(time.Duration) {}
}