orca/internal/sshpush/fanout_test.go

package sshpush

import (
	"context"
	"errors"
	"fmt"
	"sync/atomic"
	"testing"

	"golang.org/x/crypto/ssh"

	"git.cloudinit.dev/coreci/orca/internal/emitter"
)

// --- ExecAll tests ---

func TestExecAll_AllSucceed(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	// Use one server for all peers (same addr).
	addr := srv.addr()
	peers := []string{addr, addr, addr}
	out, errs := tr.ExecAll(context.Background(), peers, "echo hello")
	for _, p := range peers {
		if e, ok := errs[p]; ok && e != nil {
			t.Errorf("peer %s: %v", p, e)
		}
		if string(out[p]) != "hello\n" {
			t.Errorf("out[%s] = %q, want hello\\n", p, out[p])
		}
	}
}

func TestExecAll_OneFails(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	addr := srv.addr()
	// Peer "bad" returns a permanent error via a mock session.
	goodPeers := []string{addr}
	badPeer := "127.0.0.1:1" // unreachable -> transient dial error, retried, fails
	peers := append(goodPeers, badPeer)
	out, errs := tr.ExecAll(context.Background(), peers, "echo hello")
	if string(out[addr]) != "hello\n" {
		t.Errorf("good peer out = %q, want hello\\n", out[addr])
	}
	if errs[badPeer] == nil {
		t.Error("bad peer should have an error")
	}
}

func TestExecAll_WithConcurrency(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	addr := srv.addr()
	peers := []string{addr, addr, addr, addr}
	out, errs := tr.ExecAllWithConcurrency(context.Background(), peers, "echo hello", 2)
	for _, p := range peers {
		if e := errs[p]; e != nil {
			t.Errorf("peer %s: %v", p, e)
		}
		if string(out[p]) != "hello\n" {
			t.Errorf("out[%s] = %q", p, out[p])
		}
	}
}

func TestExecAll_EmptyPeers(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	out, errs := tr.ExecAll(context.Background(), nil, "echo hello")
	if len(out) != 0 || len(errs) != 0 {
		t.Errorf("empty peers: out=%v errs=%v", out, errs)
	}
}

func TestExecAll_ContextCancelled(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	addr := srv.addr()
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	out, errs := tr.ExecAll(ctx, []string{addr, addr}, "echo hello")
	// With a cancelled context, all peers should fail.
	for _, p := range []string{addr, addr} {
		if errs[p] == nil && string(out[p]) == "" {
			// acceptable: either error or no output
		}
	}
}

// --- WriteAll tests ---

func TestWriteAll_AllSucceed(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	addr := srv.addr()
	files := map[string][]emitter.File{
		addr: {
			{Path: "/w/a", Content: "alpha\n", Mode: "0644"},
			{Path: "/w/b", Content: "beta\n", Mode: "0644"},
		},
	}
	errs := tr.WriteAll(context.Background(), []string{addr}, files)
	for p, e := range errs {
		if e != nil {
			t.Errorf("peer %s: %v", p, e)
		}
	}
	srv.mu.Lock()
	if srv.files["/w/a"] != "alpha\n" {
		t.Errorf("file a = %q", srv.files["/w/a"])
	}
	if srv.files["/w/b"] != "beta\n" {
		t.Errorf("file b = %q", srv.files["/w/b"])
	}
	srv.mu.Unlock()
}

func TestWriteAll_OnePeerFails(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	addr := srv.addr()
	bad := "127.0.0.1:1"
	files := map[string][]emitter.File{
		addr: {{Path: "/ok/f", Content: "ok\n", Mode: "0644"}},
		bad:  {{Path: "/fail/f", Content: "fail\n", Mode: "0644"}},
	}
	errs := tr.WriteAll(context.Background(), []string{addr, bad}, files)
	if errs[addr] != nil {
		t.Errorf("good peer should not have error, got %v", errs[addr])
	}
	if errs[bad] == nil {
		t.Error("bad peer should have error")
	}
}

func TestWriteAll_EmptyPeers(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	errs := tr.WriteAll(context.Background(), nil, nil)
	if len(errs) != 0 {
		t.Errorf("empty peers: errs=%v", errs)
	}
}

func TestWriteAll_WithConcurrency(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	addr := srv.addr()
	files := map[string][]emitter.File{
		addr: {{Path: "/c/f", Content: "c\n", Mode: "0644"}},
	}
	errs := tr.WriteAllWithConcurrency(context.Background(), []string{addr}, files, 4)
	for _, e := range errs {
		if e != nil {
			t.Errorf("peer err: %v", e)
		}
	}
}

func TestWriteAll_Idempotent(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	srv.mu.Lock()
	srv.files["/i/f"] = "same\n"
	srv.mu.Unlock()
	tr := realTransport(t, srv)
	defer tr.Close()
	addr := srv.addr()
	files := map[string][]emitter.File{
		addr: {{Path: "/i/f", Content: "same\n", Mode: "0644"}},
	}
	// Capture writes to verify idempotent skip.
	var writes int64
	tr.SetSessionFactory(func(c *ssh.Client) (sshSession, error) {
		return &writeCountingSession{srv: srv, writes: &writes}, nil
	})
	// Pre-populate pool.
	client, err := tr.dial(addr)
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	tr.pool.Store(addr, client)
	errs := tr.WriteAll(context.Background(), []string{addr}, files)
	if errs[addr] != nil {
		t.Errorf("WriteAll err: %v", errs[addr])
	}
	// sha256sum returns a hash that matches -> no write command.
	srv.mu.Lock()
	content := srv.files["/i/f"]
	srv.mu.Unlock()
	if content != "same\n" {
		t.Errorf("content changed to %q", content)
	}
}

// writeCountingSession counts how many write commands (mkdir + cat >)
// are issued; returns the server's file content for sha256sum.
type writeCountingSession struct {
	srv    *fakeSSHServer
	writes *int64
}

func (w *writeCountingSession) CombinedOutput(cmd string) ([]byte, error) {
	c := trim(cmd)
	if startsWith(c, "sha256sum ") {
		path := unquote(trimPrefix(c, "sha256sum "))
		path = trimSuffix(path, " 2>/dev/null")
		w.srv.mu.Lock()
		content, ok := w.srv.files[path]
		w.srv.mu.Unlock()
		if !ok {
			return []byte(""), nil
		}
		sum := sha256HexStr([]byte(content))
		return []byte(sum + "  " + path + "\n"), nil
	}
	if startsWith(c, "mkdir -p ") && contains(c, "cat >") {
		atomic.AddInt64(w.writes, 1)
		return nil, nil
	}
	return nil, nil
}
func (w *writeCountingSession) Close() error { return nil }

// Local string helpers to avoid importing strings in a way that
// conflicts with the test's existing imports.
func trim(s string) string {
	for len(s) > 0 && (s[0] == ' ' || s[0] == '\t') {
		s = s[1:]
	}
	for len(s) > 0 && (s[len(s)-1] == ' ' || s[len(s)-1] == '\t') {
		s = s[:len(s)-1]
	}
	return s
}
func startsWith(s, prefix string) bool { return len(s) >= len(prefix) && s[:len(prefix)] == prefix }
func contains(s, sub string) bool {
	return len(sub) == 0 || (len(s) >= len(sub) && indexOf(s, sub) >= 0)
}
func indexOf(s, sub string) int {
	for i := 0; i+len(sub) <= len(s); i++ {
		if s[i:i+len(sub)] == sub {
			return i
		}
	}
	return -1
}
func trimPrefix(s, prefix string) string {
	if startsWith(s, prefix) {
		return s[len(prefix):]
	}
	return s
}
func trimSuffix(s, suffix string) string {
	if len(s) >= len(suffix) && s[len(s)-len(suffix):] == suffix {
		return s[:len(s)-len(suffix)]
	}
	return s
}

func TestDefaultFanoutConcurrency(t *testing.T) {
	if DefaultFanoutConcurrency != 8 {
		t.Errorf("DefaultFanoutConcurrency = %d, want 8", DefaultFanoutConcurrency)
	}
}

func TestParseMode_Fanout(t *testing.T) {
	// parseMode is in fanout.go; sanity-check here too.
	for _, tc := range []struct{ in, want string }{
		{"0644", "644"},
		{"0755", "755"},
		{"bad", "644"},
	} {
		got := fmt.Sprintf("%o", parseMode(tc.in))
		if got != tc.want {
			t.Errorf("parseMode(%q) = %s, want %s", tc.in, got, tc.want)
		}
	}
}

var _ = errors.New