apiary/honeypot/ssh/store/postgres_test.go

// nolint: errcheck
package store_test

import (
	"database/sql"
	"os"
	"testing"

	"git.t-juice.club/torjus/apiary/honeypot/ssh/store"
)

func TestPostgresStore(t *testing.T) {
	var dsn string
	var found bool
	dsn, found = os.LookupEnv("APIARY_TEST_POSTGRES_DSN")
	if !found {
		t.Skipf("APIARY_TEST_POSTGRES_DSN not set. Skipping.")
	}

	dropPGDatabase(dsn)

	s, err := store.NewPostgresStore(dsn)
	if err != nil {
		t.Fatalf("Error getting store: %s", err)
	}

	s.InitDB()

	testLoginAttemptStore(s, t)
}

func TestPostgresStoreWithCache(t *testing.T) {
	var dsn string
	var found bool
	dsn, found = os.LookupEnv("APIARY_TEST_POSTGRES_DSN")
	if !found {
		t.Skipf("APIARY_TEST_POSTGRES_DSN not set. Skipping.")
	}

	dropPGDatabase(dsn)

	pgs, err := store.NewPostgresStore(dsn)
	if err != nil {
		t.Fatalf("Error getting store: %s", err)
	}

	pgs.InitDB()
	s := store.NewCachingStore(pgs)

	testLoginAttemptStore(s, t)
}

func BenchmarkPostgresStore(b *testing.B) {
	var dsn string
	var found bool
	dsn, found = os.LookupEnv("APIARY_TEST_POSTGRES_DSN")
	if !found {
		b.Skipf("APIARY_TEST_POSTGRES_DSN not set. Skipping.")
	}
	dropPGDatabase(dsn)
	setupFunc := func() store.LoginAttemptStore {
		dropPGDatabase(dsn)
		pgs, err := store.NewPostgresStore(dsn)
		if err != nil {
			b.Fatalf("Error getting store: %s", err)
		}
		pgs.InitDB()
		return pgs
	}
	benchmarkLoginAttemptStore(setupFunc, b)
	dropPGDatabase(dsn)
}

func dropPGDatabase(dsn string) {
	db, err := sql.Open("pgx", dsn)
	if err != nil {
		panic(err)
	}

	_, _ = db.Exec("DROP TABLE login_attempts")
}