llzai/axonhub

1package contexts
2
3import (
4	"context"
5	"testing"
6
7	"github.com/looplj/axonhub/internal/ent"
8)
9
10func TestWithThread(t *testing.T) {
11	ctx := context.Background()
12	thread := &ent.Thread{
13		ID:       1,
14		ThreadID: "thread-123",
15	}
16
17	// Test storing thread entity
18	newCtx := WithThread(ctx, thread)
19	if newCtx == ctx {
20		t.Error("WithThread should return a new context")
21	}
22
23	// Test retrieving thread entity
24	retrievedThread, ok := GetThread(newCtx)
25	if !ok {
26		t.Error("GetThread should return true for existing thread")
27	}
28
29	if retrievedThread == nil {
30		t.Error("GetThread should return non-nil thread")
31	}
32
33	if retrievedThread.ID != thread.ID {
34		t.Errorf("expected ID %d, got %d", thread.ID, retrievedThread.ID)
35	}
36
37	if retrievedThread.ThreadID != thread.ThreadID {
38		t.Errorf("expected ThreadID %s, got %s", thread.ThreadID, retrievedThread.ThreadID)
39	}
40}
41
42func TestGetThread(t *testing.T) {
43	ctx := context.Background()
44
45	// Test retrieving thread from empty context
46	thread, ok := GetThread(ctx)
47	if ok {
48		t.Error("GetThread should return false for empty context")
49	}
50
51	if thread != nil {
52		t.Error("GetThread should return nil for empty context")
53	}
54
55	// Test retrieving thread from context with other values
56	ctxWithOtherValue := context.WithValue(ctx, "other_key", "other_value")
57
58	thread, ok = GetThread(ctxWithOtherValue)
59	if ok {
60		t.Error("GetThread should return false for context without thread")
61	}
62
63	if thread != nil {
64		t.Error("GetThread should return nil for context without thread")
65	}
66}
67
68func TestThreadWithMultipleValues(t *testing.T) {
69	ctx := context.Background()
70
71	// Test storing thread along with other values
72	ctx = WithAPIKey(ctx, &ent.APIKey{ID: 1, Key: "test-key"})
73	ctx = WithUser(ctx, &ent.User{ID: 123, Email: "test@example.com"})
74	ctx = WithThread(ctx, &ent.Thread{ID: 1, ThreadID: "thread-123"})
75	ctx = WithProjectID(ctx, 456)
76
77	// Test retrieving all values
78	apiKey, ok := GetAPIKey(ctx)
79	if !ok || apiKey.ID != 1 {
80		t.Error("API key should be stored and retrievable")
81	}
82
83	user, ok := GetUser(ctx)
84	if !ok || user.ID != 123 {
85		t.Error("User should be stored and retrievable")
86	}
87
88	thread, ok := GetThread(ctx)
89	if !ok || thread.ID != 1 {
90		t.Error("Thread should be stored and retrievable")
91	}
92
93	projectID, ok := GetProjectID(ctx)
94	if !ok || projectID != 456 {
95		t.Error("Project ID should be stored and retrievable")
96	}
97}
98
99func TestThreadOverwrite(t *testing.T) {
100	ctx := context.Background()
101
102	// Test overwriting existing thread
103	ctx = WithThread(ctx, &ent.Thread{ID: 1, ThreadID: "thread-1"})
104	ctx = WithThread(ctx, &ent.Thread{ID: 2, ThreadID: "thread-2"})
105
106	thread, ok := GetThread(ctx)
107	if !ok {
108		t.Error("Thread should exist")
109	}
110
111	if thread.ID != 2 || thread.ThreadID != "thread-2" {
112		t.Error("Thread should be the overwritten value")
113	}
114}
115