107 lines
3.0 KiB
Go
107 lines
3.0 KiB
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")
|
|
}
|
|
}
|