(t *testing.T)
| 66 | } |
| 67 | |
| 68 | func TestChatStreamRelay(t *testing.T) { |
| 69 | t.Parallel() |
| 70 | |
| 71 | t.Run("RelayMessagePartsAcrossReplicas", func(t *testing.T) { |
| 72 | t.Parallel() |
| 73 | ctx := testutil.Context(t, testutil.WaitLong) |
| 74 | |
| 75 | db, pubsub := dbtestutil.NewDB(t) |
| 76 | firstClient, firstUser := coderdenttest.New(t, &coderdenttest.Options{ |
| 77 | Options: &coderdtest.Options{ |
| 78 | Database: db, |
| 79 | Pubsub: pubsub, |
| 80 | DeploymentValues: coderdtest.DeploymentValues(t, func(dv *codersdk.DeploymentValues) { |
| 81 | require.NoError(t, dv.AI.Chat.AIGatewayRoutingEnabled.Set("false")) |
| 82 | }), |
| 83 | }, |
| 84 | LicenseOptions: &coderdenttest.LicenseOptions{ |
| 85 | Features: license.Features{ |
| 86 | codersdk.FeatureHighAvailability: 1, |
| 87 | }, |
| 88 | }, |
| 89 | }) |
| 90 | |
| 91 | secondClient, _ := coderdenttest.New(t, &coderdenttest.Options{ |
| 92 | Options: &coderdtest.Options{ |
| 93 | Database: db, |
| 94 | Pubsub: pubsub, |
| 95 | DeploymentValues: coderdtest.DeploymentValues(t, func(dv *codersdk.DeploymentValues) { |
| 96 | require.NoError(t, dv.AI.Chat.AIGatewayRoutingEnabled.Set("false")) |
| 97 | }), |
| 98 | }, |
| 99 | DontAddLicense: true, |
| 100 | DontAddFirstUser: true, |
| 101 | }) |
| 102 | secondClient.SetSessionToken(firstClient.SessionToken()) |
| 103 | |
| 104 | // Verify we have two replicas |
| 105 | replicas, err := secondClient.Replicas(ctx) |
| 106 | require.NoError(t, err) |
| 107 | require.Len(t, replicas, 2) |
| 108 | firstReplicaID := replicaIDForClientURL(t, firstClient.URL, replicas) |
| 109 | secondReplicaID := replicaIDForClientURL(t, secondClient.URL, replicas) |
| 110 | |
| 111 | streamingChunks := make(chan chattest.OpenAIChunk, 8) |
| 112 | chatStreamStarted := make(chan struct{}, 1) |
| 113 | openai := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { |
| 114 | if req.Stream { |
| 115 | select { |
| 116 | case chatStreamStarted <- struct{}{}: |
| 117 | default: |
| 118 | } |
| 119 | return chattest.OpenAIResponse{StreamingChunks: streamingChunks} |
| 120 | } |
| 121 | return chattest.OpenAINonStreamingResponse("ok") |
| 122 | }) |
| 123 | |
| 124 | expClient := codersdk.NewExperimentalClient(firstClient) |
| 125 | model := createOpenAIModelConfigForTest(ctx, t, expClient, "test", openai) |
nothing calls this directly
no test coverage detected