AetherForge/agent/client/handlemessage_test.go

package client

import (
	"encoding/json"
	"testing"
	"time"

	"crypto-miner-agent/config"
	"crypto-miner-agent/job"
	"crypto-miner-agent/miner"
	"crypto-miner-agent/stats"
)

// newTestClient builds a minimal AgentClient whose pool is running with a
// real engine so handleMessage can call pool.SetJob without panicking.
func newTestClient(t *testing.T) *AgentClient {
	t.Helper()
	b := config.GetBuiltinConfig()
	b.Threads = 1
	cfg := config.RuntimeConfig{BuiltinConfig: b}
	reporter := stats.NewReporter()
	pool := miner.NewPool(1, cfg, reporter, func(jobID, nonce, hash string) {})
	c := &AgentClient{
		cfg:  cfg,
		pool: pool,
	}
	// conn intentionally nil — write() returns error but never panics.
	return c
}

func encodePayload(t *testing.T, v interface{}) json.RawMessage {
	t.Helper()
	b, err := json.Marshal(v)
	if err != nil {
		t.Fatalf("json.Marshal: %v", err)
	}
	return json.RawMessage(b)
}

// ─── new_job dispatch ────────────────────────────────────────────────────────

func TestHandleMessageNewJobSetsJob(t *testing.T) {
	c := newTestClient(t)
	j := job.Job{
		ID:       "job-abc",
		Blob:     "0c0c9eb2e1d805a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b200000000c8d9e0f1a2b3c4d5e6f7a8b9c0d1e2f3a4b5c6d7e8f9a0b1c2d3e4f5a6b7c8d901",
		Target:   "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff",
		SeedHash: "4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a4a",
		Height:   3000000,
	}
	msg := Message{Type: "new_job", Payload: encodePayload(t, j)}

	before := time.Now().Add(-time.Millisecond)
	c.handleMessage(msg)

	// lastJobAt should be set to approximately now.
	raw := c.lastJobAt.Load()
	if raw == nil {
		t.Fatal("lastJobAt not set after new_job")
	}
	if !raw.(time.Time).After(before) {
		t.Errorf("lastJobAt=%v, want after %v", raw.(time.Time), before)
	}
}

func TestHandleMessageAISnapshotRequest(t *testing.T) {
	c := newTestClient(t)
	// Should schedule push without panicking (write fails safely with nil conn).
	c.handleMessage(Message{Type: "ai_snapshot_request", Payload: json.RawMessage("{}")})
}

func TestHandleMessageNewJobWithErrorRetries(t *testing.T) {
	c := newTestClient(t)
	// Error payload — pool not ready yet.
	payload := encodePayload(t, map[string]string{"error": "pool not connected"})
	msg := Message{Type: "new_job", Payload: payload}

	// Should not panic; AfterFunc fires after 3s but we don't wait.
	c.handleMessage(msg)

	// lastJobAt must remain unset (no valid job was delivered).
	if raw := c.lastJobAt.Load(); raw != nil {
		t.Errorf("lastJobAt should be nil on error response, got %v", raw)
	}
}

func TestHandleMessageNewJobEmptyBlobNotStored(t *testing.T) {
	c := newTestClient(t)
	j := job.Job{ID: "empty-blob", Blob: "", Target: "ff", SeedHash: "aa"}
	msg := Message{Type: "new_job", Payload: encodePayload(t, j)}

	c.handleMessage(msg)

	// Empty blob → not stored as a real job.
	if raw := c.lastJobAt.Load(); raw != nil {
		t.Error("lastJobAt should be nil for empty-blob job")
	}
}

// ─── share_result dispatch ───────────────────────────────────────────────────

func TestHandleMessageShareResultAccepted(t *testing.T) {
	c := newTestClient(t)
	result := ShareResult{Accepted: true, JobID: "job-1"}
	msg := Message{Type: "share_result", Payload: encodePayload(t, result)}

	c.handleMessage(msg)

	c.mu.Lock()
	accepted := c.sharesAccepted
	c.mu.Unlock()
	if accepted != 1 {
		t.Errorf("sharesAccepted = %d, want 1", accepted)
	}
}

func TestHandleMessageShareResultRejected(t *testing.T) {
	c := newTestClient(t)
	result := ShareResult{Accepted: false, JobID: "job-1"}
	msg := Message{Type: "share_result", Payload: encodePayload(t, result)}

	c.handleMessage(msg)

	c.mu.Lock()
	accepted := c.sharesAccepted
	c.mu.Unlock()
	if accepted != 0 {
		t.Errorf("sharesAccepted = %d, want 0 for rejected share", accepted)
	}
}

func TestHandleMessageShareResultMultiple(t *testing.T) {
	c := newTestClient(t)
	for i := 0; i < 5; i++ {
		msg := Message{Type: "share_result", Payload: encodePayload(t, ShareResult{Accepted: true})}
		c.handleMessage(msg)
	}
	c.mu.Lock()
	n := c.sharesAccepted
	c.mu.Unlock()
	if n != 5 {
		t.Errorf("sharesAccepted = %d after 5 accepted, want 5", n)
	}
}

// ─── command dispatch ─────────────────────────────────────────────────────────

func cmdPayload(t *testing.T, action string) json.RawMessage {
	return encodePayload(t, map[string]interface{}{"action": action})
}

func TestHandleMessageCommandPause(t *testing.T) {
	c := newTestClient(t)
	msg := Message{Type: "command", Payload: cmdPayload(t, "pause")}
	c.handleMessage(msg)

	// Give the goroutine a moment to run handleCommand.
	time.Sleep(10 * time.Millisecond)
	if !c.pool.IsRemotePaused() {
		t.Error("pool should be paused after pause command")
	}
}

func TestHandleMessageCommandResume(t *testing.T) {
	c := newTestClient(t)
	c.pool.PauseRemote() // pre-pause
	msg := Message{Type: "command", Payload: cmdPayload(t, "resume")}
	c.handleMessage(msg)

	time.Sleep(10 * time.Millisecond)
	if c.pool.IsRemotePaused() {
		t.Error("pool should be resumed after resume command")
	}
}

func TestHandleMessageCommandUnknownNocrash(t *testing.T) {
	c := newTestClient(t)
	msg := Message{Type: "command", Payload: cmdPayload(t, "definitely_unknown_action")}
	// Must not panic.
	c.handleMessage(msg)
	time.Sleep(10 * time.Millisecond)
}

func TestHandleMessageUnknownTypeIgnored(t *testing.T) {
	c := newTestClient(t)
	msg := Message{Type: "totally_made_up", Payload: json.RawMessage(`{}`)}
	c.handleMessage(msg) // must not panic
}

func TestHandleMessageInvalidJSON(t *testing.T) {
	c := newTestClient(t)
	msg := Message{Type: "share_result", Payload: json.RawMessage(`not json`)}
	c.handleMessage(msg) // must not panic
}

// ─── needsStratumFallback thresholds ─────────────────────────────────────────

func TestNeedsStratumFallbackOfflineUnderThreshold(t *testing.T) {
	c := newTestClient(t)
	// C2 disconnected 3 seconds ago — under the 8s threshold.
	c.connected.Store(false)
	disconnected := time.Now().Add(-3 * time.Second)
	if c.needsStratumFallback(disconnected) {
		t.Error("should not trigger fallback before 8s offline threshold")
	}
}

func TestNeedsStratumFallbackOfflineOverThreshold(t *testing.T) {
	c := newTestClient(t)
	// C2 disconnected 10 seconds ago — over the 8s threshold.
	c.connected.Store(false)
	disconnected := time.Now().Add(-10 * time.Second)
	if !c.needsStratumFallback(disconnected) {
		t.Error("should trigger fallback after 8s offline")
	}
}

func TestNeedsStratumFallbackOfflineZeroTime(t *testing.T) {
	c := newTestClient(t)
	c.connected.Store(false)
	// Zero time — disconnectedSince not yet initialized.
	if c.needsStratumFallback(time.Time{}) {
		t.Error("should not trigger fallback with zero disconnectedSince")
	}
}

func TestNeedsStratumFallbackConnectedNoJob(t *testing.T) {
	c := newTestClient(t)
	c.connected.Store(true)
	// Connected, never received a job. disconnectedSince is 20s ago.
	disconnected := time.Now().Add(-20 * time.Second)
	if !c.needsStratumFallback(disconnected) {
		t.Error("should trigger fallback when connected >15s with no job")
	}
}

func TestNeedsStratumFallbackConnectedRecentJob(t *testing.T) {
	c := newTestClient(t)
	c.connected.Store(true)
	// Received a job 5s ago — still fresh.
	c.lastJobAt.Store(time.Now().Add(-5 * time.Second))
	disconnected := time.Now().Add(-20 * time.Second)
	if c.needsStratumFallback(disconnected) {
		t.Error("should NOT trigger fallback when C2 delivered job <15s ago")
	}
}

func TestNeedsStratumFallbackConnectedStaleJob(t *testing.T) {
	c := newTestClient(t)
	c.connected.Store(true)
	// Last job was 20s ago — stale.
	c.lastJobAt.Store(time.Now().Add(-20 * time.Second))
	if !c.needsStratumFallback(time.Time{}) {
		t.Error("should trigger fallback when last job >15s ago")
	}
}