dsfree_service_test.go

package main

import (
	"bytes"
	"errors"
	"log"
	"os"
	"strings"
	"sync"
	"syscall"
	"testing"
	"time"

	"github.com/fsnotify/fsnotify"
)

type SpyHostsManager struct {
	mu                   sync.Mutex
	applyCallCount       int
	restoreCalled        bool
	failApplyOnCall      int // Fails on the Nth call to Apply
	appliedDomains       []string
	failApplyWithError   error
	failRestoreWithError error
}

func (s *SpyHostsManager) Apply(domains []string) error {
	if s.failApplyWithError != nil {
		return s.failApplyWithError
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.applyCallCount++
	s.appliedDomains = domains
	if s.failApplyOnCall > 0 && s.applyCallCount == s.failApplyOnCall {
		return errors.New("simulated apply failure on specific call")
	}
	return nil
}

func (s *SpyHostsManager) Restore() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.restoreCalled = true
	if s.failRestoreWithError != nil {
		return s.failRestoreWithError
	}
	return nil
}

// This mock gives the test full control over the file watcher's event and error channels.
type MockFileWatcher struct {
	mu         sync.Mutex
	eventsChan chan fsnotify.Event
	errorsChan chan error
	isClosed   bool
}

func NewMockFileWatcher() *MockFileWatcher {
	return &MockFileWatcher{
		eventsChan: make(chan fsnotify.Event),
		errorsChan: make(chan error),
	}
}

func (mfw *MockFileWatcher) Events() <-chan fsnotify.Event { return mfw.eventsChan }
func (mfw *MockFileWatcher) Errors() <-chan error          { return mfw.errorsChan }
func (mfw *MockFileWatcher) Add(name string) error         { return nil }
func (mfw *MockFileWatcher) Close() error {
	mfw.mu.Lock()
	defer mfw.mu.Unlock()
	// Only close the channels if they haven't been closed already to prevent a panic.
	if !mfw.isClosed {
		close(mfw.eventsChan)
		close(mfw.errorsChan)
		mfw.isClosed = true
	}
	return nil
}

type MockScheduler struct {
	mu       sync.Mutex
	callback func()
}

func (ms *MockScheduler) AfterFunc(d time.Duration, f func()) {
	ms.mu.Lock()
	defer ms.mu.Unlock()
	// Instead of waiting, we save the callback to be invoked manually by the test.
	ms.callback = f
}

func (ms *MockScheduler) Invoke() {
	ms.mu.Lock()
	cb := ms.callback
	ms.mu.Unlock()
	if cb != nil {
		cb()
	}
}


func TestDSFreeService_Run_CallsApplyOnceOnStart(t *testing.T) {
	spyManager := &SpyHostsManager{}
	mockWatcher := NewMockFileWatcher()
	domains := []string{"test.com"}
	shutdownChan := make(chan os.Signal, 1)
	mockScheduler := &MockScheduler{}
	service := NewDSFreeService(spyManager, mockWatcher, shutdownChan, mockScheduler, domains)

	// Run the service in a background goroutine so the test can interact with it.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		service.Run()
	}()

	time.Sleep(10 * time.Millisecond)
	mockWatcher.Close() // Force Run() à se terminer.
	wg.Wait()           // Attend que la goroutine soit VRAIMENT terminée.

	spyManager.mu.Lock()
	defer spyManager.mu.Unlock()

	if spyManager.applyCallCount != 1 {
		t.Errorf("Expected Apply() to be called exactly once, but it was called %d time(s)", spyManager.applyCallCount)
	}
	if len(spyManager.appliedDomains) != 1 || spyManager.appliedDomains[0] != "test.com" {
		t.Errorf("Run() passed wrong domains to Apply(). Expected ['test.com'], got %v", spyManager.appliedDomains)
	}
}

func TestDSFreeService_Run_DoesNotCallApplyWhenInitialApplyFails(t *testing.T) {
	expectedError := errors.New("simulated apply failure")
	spyManager := &SpyHostsManager{
		failApplyWithError: expectedError,
	}
	mockWatcher := NewMockFileWatcher()
	domains := []string{"test.com"}
	shutdownChan := make(chan os.Signal, 1)
	mockScheduler := &MockScheduler{}
	service := NewDSFreeService(spyManager, mockWatcher, shutdownChan, mockScheduler, domains)


	// For this test no need to call service.Run in it's own gorotine, the for select loop
	// is not reach, the error is returned before
	service.Run()

	spyManager.mu.Lock()
	defer spyManager.mu.Unlock()

	if spyManager.applyCallCount != 0 {
		t.Errorf("Expected Apply() not to be called when it returns an error, but it was called %d time(s)", spyManager.applyCallCount)
	}
}

func TestDSFreeService_Run_ReappliesOnTampering(t *testing.T) {
	spyManager := &SpyHostsManager{}
	mockWatcher := NewMockFileWatcher()
	domains := []string{"test.com"}
	shutdownChan := make(chan os.Signal, 1)
	mockScheduler := &MockScheduler{}
	service := NewDSFreeService(spyManager, mockWatcher, shutdownChan, mockScheduler, domains)

	// Run the service in a background goroutine so the test can interact with it.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		service.Run()
	}()

	time.Sleep(10 * time.Millisecond)
	mockWatcher.eventsChan <- fsnotify.Event{Name: hostsFilePath, Op: fsnotify.Write}
	time.Sleep(10 * time.Millisecond)
	mockWatcher.Close()
	wg.Wait() // Attend que Run() soit terminée avant de vérifier.

	spyManager.mu.Lock()
	defer spyManager.mu.Unlock()
	if spyManager.applyCallCount < 2 {
		t.Errorf("Expected Apply() to be called at least twice, but it was called %d time(s)", spyManager.applyCallCount)
	}
}

func TestDSFreeService_Run_HandlesWatcherError(t *testing.T) {
	spyManager := &SpyHostsManager{}
	mockWatcher := NewMockFileWatcher()
	domains := []string{"test.com"}
	shutdownChan := make(chan os.Signal, 1)
	mockScheduler := &MockScheduler{}
	service := NewDSFreeService(spyManager, mockWatcher, shutdownChan, mockScheduler, domains)

	var logBuffer bytes.Buffer
	log.SetOutput(&logBuffer)
	defer log.SetOutput(os.Stderr)

	// Run the service in a background goroutine so the test can interact with it.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		service.Run()
	}()

	time.Sleep(10 * time.Millisecond)
	mockWatcher.errorsChan <- errors.New("simulated watcher error")
	time.Sleep(10 * time.Millisecond)
	mockWatcher.Close()
	wg.Wait() // Déjà correct ici.

	if !strings.Contains(logBuffer.String(), "ERROR: File watcher error: simulated watcher error") {
		t.Errorf("Expected to see watcher error in logs, but got: %s", logBuffer.String())
	}
}

func TestDSFreeService_Run_CallsRestoreOnShutdown(t *testing.T) {
	spyManager := &SpyHostsManager{}
	mockWatcher := NewMockFileWatcher()
	shutdownChan := make(chan os.Signal, 1)
	domains := []string{"test.com"}
	mockScheduler := &MockScheduler{}
	service := NewDSFreeService(spyManager, mockWatcher, shutdownChan, mockScheduler, domains)

	// Run the service in a background goroutine so the test can interact with it.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		service.Run()
	}()

	time.Sleep(10 * time.Millisecond)
	shutdownChan <- syscall.SIGINT
	wg.Wait() // Attend que Run() soit terminée après avoir traité le signal.

	spyManager.mu.Lock()
	defer spyManager.mu.Unlock()

	if !spyManager.restoreCalled {
		t.Error("Expected Restore() to be called on shutdown, but it wasn't")
	}
}

func TestDSFreeService_Run_HandlesFailuresInLoop(t *testing.T) {
	var logBuffer bytes.Buffer
	log.SetOutput(&logBuffer)
	defer log.SetOutput(os.Stderr)

	spyManager := &SpyHostsManager{
		failApplyOnCall:      2,
		failRestoreWithError: errors.New("simulated restore failure"),
	}
	mockWatcher := NewMockFileWatcher()
	shutdownChan := make(chan os.Signal, 1)
	mockScheduler := &MockScheduler{}
	service := NewDSFreeService(spyManager, mockWatcher, shutdownChan, mockScheduler, []string{"test.com"})

	// Run the service in a background goroutine so the test can interact with it.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		service.Run()
	}()

	time.Sleep(10 * time.Millisecond)
	mockWatcher.eventsChan <- fsnotify.Event{Op: fsnotify.Write}
	time.Sleep(10 * time.Millisecond)
	mockWatcher.eventsChan <- fsnotify.Event{Op: fsnotify.Write}
	time.Sleep(10 * time.Millisecond)

	// The shutdown signal must be sent while Run() is active.
	shutdownChan <- syscall.SIGTERM

	wg.Wait() // wait Run terminate

	spyManager.mu.Lock()
	if spyManager.applyCallCount != 2 {
		t.Errorf("Expected Apply to be called twice, got %d", spyManager.applyCallCount)
	}
	if !spyManager.restoreCalled {
		t.Error("Expected Restore to be called, but it wasn't")
	}
	spyManager.mu.Unlock()

	logs := logBuffer.String()
	if !strings.Contains(logs, "Failed to re-apply sabotage") {
		t.Error("Expected to see log for failed re-apply")
	}
	if !strings.Contains(logs, "Could not restore hosts file on shutdown") {
		t.Error("Expected to see log for failed restore")
	}
}

func TestDSFreeService_Run_ResetsCooldownAfterTampering(t *testing.T) {
	spyManager := &SpyHostsManager{}
	mockWatcher := NewMockFileWatcher()
	shutdownChan := make(chan os.Signal, 1)
	mockScheduler := &MockScheduler{}
	domains := []string{"test.com"}
	service := NewDSFreeService(spyManager, mockWatcher, shutdownChan, mockScheduler, domains)

	// Run the service in a background goroutine so the test can interact with it.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		service.Run()
	}()

	time.Sleep(10 * time.Millisecond)
	mockWatcher.eventsChan <- fsnotify.Event{Op: fsnotify.Write}
	time.Sleep(10 * time.Millisecond)

	service.mu.Lock()
	if !service.isCoolingDown {
		service.mu.Unlock()
		t.Fatal("Expected cooldown to be active immediately after tampering event")
	}
	service.mu.Unlock()

	mockScheduler.Invoke()

	service.mu.Lock()
	if service.isCoolingDown {
		service.mu.Unlock()
		t.Fatal("Expected cooldown to be false after scheduler callback")
	}
	service.mu.Unlock()

	mockWatcher.Close()
	wg.Wait()
}