Smail/backend/services/dbmigration/maintenance_test.go

package dbmigration

import (
	"context"
	"database/sql"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func TestDatabaseDrainWaitsForBorrowedConnectionAfterPoolClose(t *testing.T) {
	database, err := sql.Open("sqlite", ":memory:")
	if err != nil {
		t.Fatal(err)
	}
	defer database.Close()
	connection, err := database.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer connection.Close()
	if err := database.Close(); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
	defer cancel()
	if err := WaitForDatabaseDrain(ctx, database); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("borrowed connection was not protected: %v", err)
	}
	if err := connection.Close(); err != nil {
		t.Fatal(err)
	}
	if err := WaitForDatabaseDrain(context.Background(), database); err != nil {
		t.Fatal(err)
	}
}

func TestMaintenanceGuardDrainsRequestsAndKeepsRestartRoute(t *testing.T) {
	executor := &PortableExecutor{}
	service := NewService(&DatabaseExecutor{portable: executor})
	previous := Current()
	SetCurrent(service)
	t.Cleanup(func() { SetCurrent(previous) })
	entered, release, migrated := make(chan struct{}), make(chan struct{}), make(chan struct{})
	handler := GuardHTTP(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/write":
			close(entered)
			<-release
		case "/api/v1/database/migrations":
			release, err := service.drainRequests(r.Context())
			if err != nil {
				http.Error(w, err.Error(), http.StatusRequestTimeout)
				return
			}
			defer release()
			executor.paused.Store(true)
			close(migrated)
		}
		w.WriteHeader(http.StatusOK)
	}))
	done := make(chan struct{})
	go func() {
		handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/write", nil))
		close(done)
	}()
	<-entered
	migrationDone := make(chan struct{})
	go func() {
		handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/api/v1/database/migrations", nil))
		close(migrationDone)
	}()
	select {
	case <-migrated:
		t.Error("migration did not drain request")
	case <-time.After(20 * time.Millisecond):
	}
	close(release)
	<-done
	<-migrationDone
	for _, test := range []struct {
		method, path string
		code         int
	}{
		{http.MethodPost, "/api/v1/email/send", 503},
		{http.MethodPost, "/api/v1/database/migrations", 503},
		{http.MethodPut, "/api/v1/config/base", 503},
		{http.MethodGet, "/api/v1/config/status", 200},
		{http.MethodGet, "/api/v1/config/base", 200},
		{http.MethodPost, "/api/v1/user/info", 200},
		{http.MethodPost, "/api/v1/system/restart", 200},
		{http.MethodPost, "/api/v1/system/restart/preview", 200},
	} {
		recorder := httptest.NewRecorder()
		handler.ServeHTTP(recorder, httptest.NewRequest(test.method, test.path, nil))
		if recorder.Code != test.code {
			t.Errorf("%s: %d", test.path, recorder.Code)
		}
	}
	called := false
	if err := WhileWritable(func() error { called = true; return nil }); err != nil || called {
		t.Fatal("cleanup wrote to frozen source")
	}
}

func TestMaintenanceDrainsBackgroundDeliveryBeforeFreeze(t *testing.T) {
	executor := &PortableExecutor{}
	service := NewService(&DatabaseExecutor{portable: executor})
	previous := Current()
	SetCurrent(service)
	t.Cleanup(func() { SetCurrent(previous) })
	entered, finish, done := make(chan struct{}), make(chan struct{}), make(chan struct{})
	go func() {
		defer close(done)
		_ = WhileWritable(func() error {
			close(entered)
			<-finish
			return nil
		})
	}()
	<-entered
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	release, err := service.drainRequests(ctx)
	release()
	close(finish)
	<-done
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("migration did not wait for delivery: %v", err)
	}
	release, err = service.drainRequests(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	executor.paused.Store(true)
	release()
	called := false
	if err := WhileWritable(func() error { called = true; return nil }); err != nil || called {
		t.Fatal("delivery started after source freeze")
	}
}