main pico / pkg / apps / prose / uploader_test.go
Antonio Mika  ·  2026-10-11
  1package prose
  2
  3import (
  4	"context"
  5	"fmt"
  6	"log/slog"
  7	"slices"
  8	"strings"
  9	"testing"
 10	"time"
 11
 12	"github.com/picosh/pico/pkg/db"
 13	"github.com/picosh/pico/pkg/db/stub"
 14	"github.com/picosh/pico/pkg/pssh"
 15	sendutils "github.com/picosh/pico/pkg/send/utils"
 16	"github.com/picosh/pico/pkg/shared"
 17	"golang.org/x/crypto/ssh"
 18)
 19
 20type mockFeatureDB struct {
 21	*stub.StubDB
 22	features map[string]*db.FeatureFlag
 23}
 24
 25func newMockFeatureDB() *mockFeatureDB {
 26	return &mockFeatureDB{
 27		StubDB:   stub.NewStubDB(slog.Default()),
 28		features: make(map[string]*db.FeatureFlag),
 29	}
 30}
 31
 32func (m *mockFeatureDB) FindFeature(userID, name string) (*db.FeatureFlag, error) {
 33	key := userID + ":" + name
 34	if ff, ok := m.features[key]; ok {
 35		return ff, nil
 36	}
 37	return nil, fmt.Errorf("feature flag %s not found for user %s", name, userID)
 38}
 39
 40func (m *mockFeatureDB) setFeature(userID string, ff *db.FeatureFlag) {
 41	key := userID + ":" + ff.Name
 42	m.features[key] = ff
 43}
 44
 45func TestUploadHandlerValidate(t *testing.T) {
 46	cfg := NewConfig("prose-test")
 47	validExpires := time.Now().Add(24 * time.Hour)
 48	userID := "user-1"
 49
 50	createSession := func(user *db.User) *pssh.SSHServerConnSession {
 51		conn := &pssh.SSHServerConn{
 52			Conn: &ssh.ServerConn{
 53				Permissions: &ssh.Permissions{
 54					Extensions: map[string]string{},
 55				},
 56			},
 57			Logger: slog.Default(),
 58		}
 59		sesh := &pssh.SSHServerConnSession{
 60			SSHServerConn: conn,
 61			Ctx:           context.Background(),
 62		}
 63		if user != nil {
 64			pssh.SetUser(sesh, user)
 65		}
 66		return sesh
 67	}
 68
 69	t.Run("nil user in session fails validation", func(t *testing.T) {
 70		mockDB := newMockFeatureDB()
 71		handler := NewUploadHandler(cfg, mockDB, map[string]ReadWriteHandler{})
 72
 73		sesh := createSession(nil)
 74		err := handler.Validate(sesh)
 75		if err == nil {
 76			t.Fatalf("expected error for nil user, got nil")
 77		}
 78	})
 79
 80	t.Run("user without feature flag fails validation", func(t *testing.T) {
 81		mockDB := newMockFeatureDB()
 82		handler := NewUploadHandler(cfg, mockDB, map[string]ReadWriteHandler{})
 83
 84		user := &db.User{ID: userID, Name: "tester"}
 85		sesh := createSession(user)
 86		err := handler.Validate(sesh)
 87		if err == nil {
 88			t.Fatalf("expected error for user without feature flag, got nil")
 89		}
 90	})
 91
 92	t.Run("user with prose feature flag passes validation and sets session flag", func(t *testing.T) {
 93		mockDB := newMockFeatureDB()
 94		proseFF := db.NewFeatureFlag(userID, "prose", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
 95		proseFF.ExpiresAt = &validExpires
 96		mockDB.setFeature(userID, proseFF)
 97
 98		handler := NewUploadHandler(cfg, mockDB, map[string]ReadWriteHandler{})
 99
100		user := &db.User{ID: userID, Name: "tester"}
101		sesh := createSession(user)
102		err := handler.Validate(sesh)
103		if err != nil {
104			t.Fatalf("unexpected error: %v", err)
105		}
106
107		ff := getFeatureFlag(sesh)
108		if ff == nil || ff.Name != "prose" {
109			t.Errorf("expected session feature flag 'prose', got %v", ff)
110		}
111	})
112}
113
114func TestPostWriteTagsAndAliases(t *testing.T) {
115	logger := slog.New(slog.DiscardHandler)
116	session := &pssh.SSHServerConnSession{
117		SSHServerConn: &pssh.SSHServerConn{
118			Conn:   &ssh.ServerConn{Permissions: &ssh.Permissions{Extensions: map[string]string{}}},
119			Logger: logger,
120		},
121		Ctx: context.Background(),
122	}
123	user := &db.User{ID: "user-1", Name: "tester"}
124	pssh.SetUser(session, user)
125	posts := stub.NewPostsDB(logger, user)
126	cfg := NewConfig("prose-test")
127	h := &postUploader{Cfg: cfg, DB: posts}
128	write := func(text string, mtime int64) *db.Post {
129		t.Helper()
130		_, err := h.Write(session, &sendutils.FileEntry{Filepath: "/hello.md", Reader: strings.NewReader(text), Mtime: mtime})
131		if err != nil {
132			t.Fatal(err)
133		}
134		post, err := posts.FindPostWithFilename("hello.md", user.ID, cfg.Space)
135		if err != nil {
136			t.Fatal(err)
137		}
138		return post
139	}
140
141	post := write("---\ntitle: Hello\ntags: [a, b]\naliases: [hi]\n---\n# hello\n", 100)
142	if post.Title != "Hello" || !slices.Equal(post.Tags, []string{"a", "b"}) || !slices.Equal(posts.Aliases(post.ID), []string{"hi"}) {
143		t.Errorf("new post = %q tags %q aliases %q", post.Title, post.Tags, posts.Aliases(post.ID))
144	}
145
146	post = write("# hello again\n", 200)
147	if len(post.Tags) != 0 || len(posts.Aliases(post.ID)) != 0 {
148		t.Errorf("updated post kept tags %q and aliases %q", post.Tags, posts.Aliases(post.ID))
149	}
150}