Propagate prompt (#5033)

This commit is contained in:
owenofbrien
2025-08-01 14:37:56 -05:00
committed by GitHub
parent 67d16992cf
commit a6a386f72a
11 changed files with 245 additions and 102 deletions

View File

@@ -14,13 +14,25 @@ vi.mock('google-auth-library');
describe('CodeAssistServer', () => {
it('should be able to be constructed', () => {
const auth = new OAuth2Client();
const server = new CodeAssistServer(auth, 'test-project');
const server = new CodeAssistServer(
auth,
'test-project',
{},
'test-session',
UserTierId.FREE,
);
expect(server).toBeInstanceOf(CodeAssistServer);
});
it('should call the generateContent endpoint', async () => {
const client = new OAuth2Client();
const server = new CodeAssistServer(client, 'test-project');
const server = new CodeAssistServer(
client,
'test-project',
{},
'test-session',
UserTierId.FREE,
);
const mockResponse = {
response: {
candidates: [
@@ -38,10 +50,13 @@ describe('CodeAssistServer', () => {
};
vi.spyOn(server, 'requestPost').mockResolvedValue(mockResponse);
const response = await server.generateContent({
model: 'test-model',
contents: [{ role: 'user', parts: [{ text: 'request' }] }],
});
const response = await server.generateContent(
{
model: 'test-model',
contents: [{ role: 'user', parts: [{ text: 'request' }] }],
},
'user-prompt-id',
);
expect(server.requestPost).toHaveBeenCalledWith(
'generateContent',
@@ -55,7 +70,13 @@ describe('CodeAssistServer', () => {
it('should call the generateContentStream endpoint', async () => {
const client = new OAuth2Client();
const server = new CodeAssistServer(client, 'test-project');
const server = new CodeAssistServer(
client,
'test-project',
{},
'test-session',
UserTierId.FREE,
);
const mockResponse = (async function* () {
yield {
response: {
@@ -75,10 +96,13 @@ describe('CodeAssistServer', () => {
})();
vi.spyOn(server, 'requestStreamingPost').mockResolvedValue(mockResponse);
const stream = await server.generateContentStream({
model: 'test-model',
contents: [{ role: 'user', parts: [{ text: 'request' }] }],
});
const stream = await server.generateContentStream(
{
model: 'test-model',
contents: [{ role: 'user', parts: [{ text: 'request' }] }],
},
'user-prompt-id',
);
for await (const res of stream) {
expect(server.requestStreamingPost).toHaveBeenCalledWith(
@@ -92,7 +116,13 @@ describe('CodeAssistServer', () => {
it('should call the onboardUser endpoint', async () => {
const client = new OAuth2Client();
const server = new CodeAssistServer(client, 'test-project');
const server = new CodeAssistServer(
client,
'test-project',
{},
'test-session',
UserTierId.FREE,
);
const mockResponse = {
name: 'operations/123',
done: true,
@@ -114,7 +144,13 @@ describe('CodeAssistServer', () => {
it('should call the loadCodeAssist endpoint', async () => {
const client = new OAuth2Client();
const server = new CodeAssistServer(client, 'test-project');
const server = new CodeAssistServer(
client,
'test-project',
{},
'test-session',
UserTierId.FREE,
);
const mockResponse = {
currentTier: {
id: UserTierId.FREE,
@@ -140,7 +176,13 @@ describe('CodeAssistServer', () => {
it('should return 0 for countTokens', async () => {
const client = new OAuth2Client();
const server = new CodeAssistServer(client, 'test-project');
const server = new CodeAssistServer(
client,
'test-project',
{},
'test-session',
UserTierId.FREE,
);
const mockResponse = {
totalTokens: 100,
};
@@ -155,7 +197,13 @@ describe('CodeAssistServer', () => {
it('should throw an error for embedContent', async () => {
const client = new OAuth2Client();
const server = new CodeAssistServer(client, 'test-project');
const server = new CodeAssistServer(
client,
'test-project',
{},
'test-session',
UserTierId.FREE,
);
await expect(
server.embedContent({
model: 'test-model',