131 lines
3.7 KiB
Go
131 lines
3.7 KiB
Go
package server
|
|
|
|
import (
|
|
"database/sql"
|
|
"encoding/json"
|
|
"os"
|
|
"testing"
|
|
|
|
"github.com/x-grk/longpolling-remote-access/internal/proto"
|
|
)
|
|
|
|
// On a real PostgreSQL (LPRA_TEST_PG: {"host","port","name","user","password",
|
|
// "sslmode"} of an empty database): the first start copies the SQLite data
|
|
// over, and the panel's queries work there.
|
|
func TestPostgres(t *testing.T) {
|
|
raw := os.Getenv("LPRA_TEST_PG")
|
|
if raw == "" {
|
|
t.Skip("LPRA_TEST_PG not set")
|
|
}
|
|
var c PGConf
|
|
if err := json.Unmarshal([]byte(raw), &c); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
// SQLite with a user, a machine, a rule, a saved host and some talk.
|
|
s, _ := newTestServer(t)
|
|
a, _ := s.db.AdminByUsername("adm")
|
|
if _, err := s.db.CreateForwardRule(ForwardRule{AgentID: "x1", ClientPort: 80, HostPort: 18080, ListenAddr: "0.0.0.0", Proto: "tcp", NoUfw: true}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := s.saveSSHHost(a.ID, savedHost{Label: "web", Host: "h", Port: 22, User: "root", Auth: "password"}, sshSecret{Password: "pw"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
chat := s.aiNewChat(a.ID)
|
|
s.aiAppendTo(a.ID, chat, "user", "привет")
|
|
_ = s.db.SetUserAgents(a.ID, []string{"x1"})
|
|
|
|
pgReset(t, c)
|
|
pg, err := OpenPostgres(c, s.db.path)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer pg.Close()
|
|
// Clean up for the next run.
|
|
defer func() {
|
|
for i := len(copyTables) - 1; i >= 0; i-- {
|
|
_, _ = pg.sql.Exec(`DROP TABLE IF EXISTS ` + copyTables[i] + ` CASCADE`)
|
|
}
|
|
}()
|
|
p := &Server{db: pg}
|
|
if got, _ := pg.AdminByUsername("adm"); got == nil || got.ID != a.ID {
|
|
t.Fatalf("admin not copied: %+v", got)
|
|
}
|
|
if ag, _ := pg.GetAgent("x1"); ag == nil {
|
|
t.Fatal("agent not copied")
|
|
}
|
|
rules, _ := pg.ListForwardRules()
|
|
if len(rules) != 1 || !rules[0].NoUfw {
|
|
t.Fatalf("rules: %+v", rules)
|
|
}
|
|
tg := &proto.SSHTarget{Saved: p.listSSHHosts(a.ID)[0].ID}
|
|
if err := p.fillSaved(a.ID, tg); err != nil || tg.Password != "pw" {
|
|
t.Fatalf("saved host: %+v %v", tg, err)
|
|
}
|
|
if h := p.aiHistoryOf(a.ID, chat, 10); len(h) != 1 || h[0].Content != "привет" {
|
|
t.Fatalf("talk: %+v", h)
|
|
}
|
|
// New rows go on after the copied ids.
|
|
id, err := pg.CreateUser("bob", "x", RoleUser)
|
|
if err != nil || id <= a.ID {
|
|
t.Fatalf("new user id %d, %v", id, err)
|
|
}
|
|
if err := pg.SetConfig("k", "v1"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := pg.SetConfig("k", "v2"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if v, _ := pg.GetConfig("k"); v != "v2" {
|
|
t.Fatalf("config upsert: %q", v)
|
|
}
|
|
if err := pg.SetUserAgents(id, []string{"x1", "x1"}); err != nil {
|
|
t.Fatalf("grants: %v", err)
|
|
}
|
|
p.aiAppendTo(a.ID, chat, "assistant", "ответ")
|
|
if h := p.aiHistoryOf(a.ID, chat, 10); len(h) != 2 {
|
|
t.Fatalf("talk after: %+v", h)
|
|
}
|
|
}
|
|
|
|
// Chats on PostgreSQL: past the limit the oldest go.
|
|
func TestPostgresChats(t *testing.T) {
|
|
raw := os.Getenv("LPRA_TEST_PG")
|
|
if raw == "" {
|
|
t.Skip("LPRA_TEST_PG not set")
|
|
}
|
|
var c PGConf
|
|
_ = json.Unmarshal([]byte(raw), &c)
|
|
s, _ := newTestServer(t)
|
|
pgReset(t, c)
|
|
pg, err := OpenPostgres(c, s.db.path)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer pg.Close()
|
|
defer func() {
|
|
for i := len(copyTables) - 1; i >= 0; i-- {
|
|
_, _ = pg.sql.Exec(`DROP TABLE IF EXISTS ` + copyTables[i] + ` CASCADE`)
|
|
}
|
|
}()
|
|
p := &Server{db: pg}
|
|
a, _ := pg.AdminByUsername("adm")
|
|
for i := 0; i < aiChatsKeep+3; i++ {
|
|
p.aiAppendTo(a.ID, p.aiNewChat(a.ID), "user", "q")
|
|
}
|
|
if n := len(p.aiChats(a.ID)); n != aiChatsKeep {
|
|
t.Fatalf("chats kept: %d", n)
|
|
}
|
|
}
|
|
|
|
// pgReset empties the test database: the copy runs only into an empty one.
|
|
func pgReset(t *testing.T, c PGConf) {
|
|
d, err := sql.Open("pgx", c.DSN())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer d.Close()
|
|
for i := len(copyTables) - 1; i >= 0; i-- {
|
|
_, _ = d.Exec(`DROP TABLE IF EXISTS ` + copyTables[i] + ` CASCADE`)
|
|
}
|
|
}
|