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}