Skip to content

Commit 708626d

Browse files
gagikgagik
committed
chore: simplify structure, improve tests, add height: 100% to portal
1 parent f65bcf6 commit 708626d

File tree

9 files changed

+297
-125
lines changed

9 files changed

+297
-125
lines changed

‎packages/compass-assistant/src/assistant-chat.spec.tsx‎

Lines changed: 15 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,17 @@ import { render, screen, userEvent } from '@mongodb-js/testing-library-compass';
33
import { AssistantChat } from './assistant-chat';
44
import { expect } from 'chai';
55
import type { UIMessage } from './@ai-sdk/react/use-chat';
6-
import { Chat } from './@ai-sdk/react/chat-react';
7-
import sinon from 'sinon';
6+
import { createMockChat } from '../test/utils';
87

98
describe('AssistantChat', function () {
109
const mockMessages: UIMessage[] = [
1110
{
12-
id: '1',
11+
id: 'user',
1312
role: 'user',
1413
parts: [{ type: 'text', text: 'Hello, MongoDB Assistant!' }],
1514
},
1615
{
17-
id: '2',
16+
id: 'assistant',
1817
role: 'assistant',
1918
parts: [
2019
{
@@ -25,27 +24,13 @@ describe('AssistantChat', function () {
2524
},
2625
];
2726

28-
let renderWithChat: (messages: UIMessage[]) => {
29-
result: ReturnType<typeof render>;
30-
chat: Chat<UIMessage> & {
31-
sendMessage: sinon.SinonStub;
27+
function renderWithChat(messages: UIMessage[]) {
28+
const chat = createMockChat({ messages });
29+
return {
30+
result: render(<AssistantChat chat={chat} />),
31+
chat,
3232
};
33-
};
34-
35-
beforeEach(() => {
36-
renderWithChat = (messages: UIMessage[]) => {
37-
const newChat = new Chat<UIMessage>({
38-
messages,
39-
});
40-
sinon.replace(newChat, 'sendMessage', sinon.stub());
41-
return {
42-
result: render(<AssistantChat chat={newChat} />),
43-
chat: newChat as unknown as Chat<UIMessage> & {
44-
sendMessage: sinon.SinonStub;
45-
},
46-
};
47-
};
48-
});
33+
}
4934

5035
it('renders input field and send button', function () {
5136
renderWithChat([]);
@@ -114,9 +99,12 @@ describe('AssistantChat', function () {
11499

115100
expect(screen.getByTestId('assistant-message-user')).to.exist;
116101
expect(screen.getByTestId('assistant-message-assistant')).to.exist;
117-
expect(screen.getByText('Hello, MongoDB Assistant!')).to.exist;
118-
expect(screen.getByText('Hello! How can I help you with MongoDB today?')).to
119-
.exist;
102+
expect(screen.getByTestId('assistant-message-user')).to.have.text(
103+
'Hello, MongoDB Assistant!'
104+
);
105+
expect(screen.getByTestId('assistant-message-assistant')).to.have.text(
106+
'Hello! How can I help you with MongoDB today?'
107+
);
120108
});
121109

122110
it('calls sendMessage when form is submitted', function () {

‎packages/compass-assistant/src/assistant-chat.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
5656
{messages.map((message) => (
5757
<div
5858
key={message.id}
59-
data-testid={`assistant-message-${message.role}`}
59+
data-testid={`assistant-message-${message.id}`}
6060
style={{
6161
marginBottom: '12px',
6262
padding: '8px 12px',

‎packages/compass-assistant/src/compass-assistant-drawer.tsx‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,34 @@
11
import React, { useContext } from 'react';
22
import { DrawerSection } from '@mongodb-js/compass-components';
33
import { AssistantChat } from './assistant-chat';
4-
import { ASSISTANT_DRAWER_ID, AssistantContext } from './assistant-provider';
4+
import {
5+
ASSISTANT_DRAWER_ID,
6+
AssistantContext,
7+
} from './compass-assistant-provider';
8+
import { usePreference } from 'compass-preferences-model/provider';
59

610
/**
711
* CompassAssistantDrawer component that wraps AssistantChat in a DrawerSection.
812
* This component can be placed at any level in the component tree as long as
913
* it's within an AssistantProvider.
1014
*/
11-
export const CompassAssistantDrawer: React.FunctionComponent = () => {
15+
export const CompassAssistantDrawer: React.FunctionComponent<{
16+
autoOpen?: boolean;
17+
}> = ({ autoOpen }) => {
1218
const context = useContext(AssistantContext);
1319

20+
const enableAIAssistant = usePreference('enableAIAssistant');
21+
22+
if (!enableAIAssistant) {
23+
return null;
24+
}
25+
1426
if (!context) {
1527
throw new Error(
1628
'CompassAssistantDrawer must be used within an CompassAssistantProvider'
1729
);
1830
}
1931

20-
if (!context.isEnabled) {
21-
return null;
22-
}
23-
2432
const { chat } = context;
2533

2634
return (
@@ -29,6 +37,7 @@ export const CompassAssistantDrawer: React.FunctionComponent = () => {
2937
title="MongoDB Assistant"
3038
label="MongoDB Assistant"
3139
glyph="Sparkle"
40+
autoOpen={autoOpen}
3241
>
3342
<AssistantChat chat={chat} />
3443
</DrawerSection>

‎packages/compass-assistant/src/compass-assistant-provider.spec.tsx‎

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
import React from 'react';
2+
import {
3+
render,
4+
screen,
5+
userEvent,
6+
waitFor,
7+
} from '@mongodb-js/testing-library-compass';
8+
import {
9+
AssistantProvider,
10+
CompassAssistantProvider,
11+
} from './compass-assistant-provider';
12+
import { expect } from 'chai';
13+
import sinon from 'sinon';
14+
import type { UIMessage } from './@ai-sdk/react/use-chat';
15+
import { Chat } from './@ai-sdk/react/chat-react';
16+
17+
import {
18+
DrawerAnchor,
19+
DrawerContentProvider,
20+
} from '@mongodb-js/compass-components';
21+
import type { AtlasService } from '@mongodb-js/atlas-service/provider';
22+
import { CompassAssistantDrawer } from './compass-assistant-drawer';
23+
import { createMockChat } from '../test/utils';
24+
25+
// Test component that renders AssistantProvider with children
26+
const TestComponent: React.FunctionComponent<{
27+
chat: Chat<UIMessage>;
28+
autoOpen?: boolean;
29+
}> = ({ chat, autoOpen }) => {
30+
return (
31+
<DrawerContentProvider>
32+
<AssistantProvider chat={chat}>
33+
<DrawerAnchor>
34+
<div data-testid="provider-children">Provider children</div>
35+
<CompassAssistantDrawer autoOpen={autoOpen} />
36+
</DrawerAnchor>
37+
</AssistantProvider>
38+
</DrawerContentProvider>
39+
);
40+
};
41+
42+
describe('AssistantProvider', function () {
43+
it('always renders children', function () {
44+
render(<TestComponent chat={createMockChat({ messages: [] })} />, {
45+
preferences: { enableAIAssistant: true },
46+
});
47+
48+
expect(screen.getByTestId('provider-children')).to.exist;
49+
});
50+
51+
it('does not render assistant drawer when AI assistant is disabled', function () {
52+
render(<TestComponent chat={createMockChat({ messages: [] })} />, {
53+
preferences: { enableAIAssistant: false },
54+
});
55+
56+
expect(screen.getByTestId('provider-children')).to.exist;
57+
// The drawer toolbar button should not exist when disabled
58+
expect(screen.queryByLabelText('MongoDB Assistant')).to.not.exist;
59+
});
60+
61+
it('renders the assistant drawer as the first drawer item when AI assistant is enabled', function () {
62+
render(<TestComponent chat={createMockChat({ messages: [] })} />, {
63+
preferences: { enableAIAssistant: true },
64+
});
65+
66+
expect(screen.getByTestId('lg-drawer-toolbar-icon_button-0')).to.have.attr(
67+
'aria-label',
68+
'MongoDB Assistant'
69+
);
70+
});
71+
72+
describe('with existing chat instance', function () {
73+
async function renderOpenAssistantDrawer(
74+
mockChat: Chat<UIMessage>
75+
): Promise<ReturnType<typeof render>> {
76+
const result = render(<TestComponent chat={mockChat} autoOpen={true} />, {
77+
preferences: { enableAIAssistant: true },
78+
});
79+
80+
await waitFor(() => {
81+
expect(screen.getByTestId('assistant-chat')).to.exist;
82+
});
83+
84+
return result;
85+
}
86+
87+
it('displays messages in the chat feed', async function () {
88+
const mockMessages: UIMessage[] = [
89+
{
90+
id: '1',
91+
role: 'user',
92+
parts: [{ type: 'text', text: 'Test message' }],
93+
},
94+
{
95+
id: '2',
96+
role: 'assistant',
97+
parts: [{ type: 'text', text: 'Test assistant message' }],
98+
},
99+
];
100+
const mockChat = createMockChat({ messages: mockMessages });
101+
102+
await renderOpenAssistantDrawer(mockChat);
103+
104+
expect(screen.getByTestId('assistant-message-1')).to.exist;
105+
expect(screen.getByTestId('assistant-message-1')).to.have.text(
106+
'Test message'
107+
);
108+
109+
expect(screen.getByTestId('assistant-message-2')).to.exist;
110+
expect(screen.getByTestId('assistant-message-2')).to.have.text(
111+
'Test assistant message'
112+
);
113+
});
114+
115+
it('handles message sending with custom chat when drawer is open', async function () {
116+
const mockChat = new Chat<UIMessage>({
117+
messages: [
118+
{
119+
id: 'assistant',
120+
role: 'assistant',
121+
parts: [{ type: 'text', text: 'Hello user!' }],
122+
},
123+
],
124+
});
125+
126+
const sendMessageSpy = sinon.spy(mockChat, 'sendMessage');
127+
128+
await renderOpenAssistantDrawer(mockChat);
129+
130+
const input = screen.getByTestId('assistant-chat-input');
131+
const sendButton = screen.getByTestId('assistant-chat-send-button');
132+
133+
userEvent.type(input, 'Hello assistant');
134+
userEvent.click(sendButton);
135+
136+
expect(sendMessageSpy.calledOnce).to.be.true;
137+
await waitFor(() => {
138+
expect(sendMessageSpy.firstCall.args[0]).to.deep.include({
139+
text: 'Hello assistant',
140+
});
141+
});
142+
});
143+
144+
it('new messages are added to the chat feed when the send button is clicked', async function () {
145+
const mockChat = new Chat<UIMessage>({
146+
messages: [
147+
{
148+
id: 'assistant',
149+
role: 'assistant',
150+
parts: [{ type: 'text', text: 'Hello user!' }],
151+
},
152+
],
153+
transport: {
154+
sendMessages: sinon.stub().returns(
155+
new Promise(() => {
156+
return new ReadableStream({});
157+
})
158+
),
159+
reconnectToStream: sinon.stub(),
160+
},
161+
});
162+
163+
const sendMessageSpy = sinon.spy(mockChat, 'sendMessage');
164+
165+
await renderOpenAssistantDrawer(mockChat);
166+
167+
userEvent.type(
168+
screen.getByTestId('assistant-chat-input'),
169+
'Hello assistant!'
170+
);
171+
userEvent.click(screen.getByTestId('assistant-chat-send-button'));
172+
173+
expect(sendMessageSpy.calledOnce).to.be.true;
174+
expect(sendMessageSpy.firstCall.args[0]).to.deep.include({
175+
text: 'Hello assistant!',
176+
});
177+
178+
await waitFor(() => {
179+
expect(screen.getByText('Hello assistant!')).to.exist;
180+
});
181+
});
182+
});
183+
184+
describe('CompassAssistantProvider', function () {
185+
beforeEach(function () {
186+
process.env.COMPASS_ASSISTANT_USE_ATLAS_SERVICE_URL = 'true';
187+
});
188+
189+
afterEach(function () {
190+
delete process.env.COMPASS_ASSISTANT_USE_ATLAS_SERVICE_URL;
191+
});
192+
193+
it('uses the Atlas Service assistantApiEndpoint', function () {
194+
const mockAtlasService = {
195+
assistantApiEndpoint: sinon
196+
.stub()
197+
.returns('https://example.com/assistant/api/v1'),
198+
};
199+
200+
const MockedProvider = CompassAssistantProvider.withMockServices({
201+
atlasService: mockAtlasService as unknown as AtlasService,
202+
});
203+
204+
render(
205+
<DrawerContentProvider>
206+
<DrawerAnchor />
207+
<MockedProvider chat={new Chat({})} />
208+
</DrawerContentProvider>,
209+
{
210+
preferences: { enableAIAssistant: true },
211+
}
212+
);
213+
214+
expect(mockAtlasService.assistantApiEndpoint.calledOnce).to.be.true;
215+
});
216+
});
217+
});

0 commit comments

Comments
 (0)