Skip to content

Commit

Permalink
[Obs AI Assistant] Update serverless tests to use cookie auth
Browse files Browse the repository at this point in the history
  • Loading branch information
@viduni94
viduni94 committed Dec 6, 2024
1 parent 8b28728 commit e4b0b80
Show file tree
Hide file tree
Showing 13 changed files with 118 additions and 215 deletions.
65 changes: 50 additions & 15 deletions ...on/test_suites/observability/ai_assistant/common/observability_ai_assistant_api_client.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,37 +15,47 @@ import supertest from 'supertest';
import { Subtract } from 'utility-types';
import { format } from 'url';
import { Config } from '@kbn/test';
import { SupertestWithRoleScope } from '@kbn/test-suites-xpack/api_integration/deployment_agnostic/services/role_scoped_supertest';
import { InheritedFtrProviderContext } from '../../../../services';
import type { InternalRequestHeader, RoleCredentials } from '../../../../../shared/services';

export function getObservabilityAIAssistantApiClient({
svlSharedConfig,
supertestUserWithCookieCredentials,
}: {
svlSharedConfig: Config;
supertestUserWithCookieCredentials?: SupertestWithRoleScope;
}) {
const kibanaServer = svlSharedConfig.get('servers.kibana');
const cAuthorities = svlSharedConfig.get('servers.kibana.certificateAuthorities');

const url = format({
...kibanaServer,
auth: false, // don't use auth in serverless
});

return createObservabilityAIAssistantApiClient(supertest.agent(url, { ca: cAuthorities }));
if (supertestUserWithCookieCredentials) {
return createObservabilityAIAssistantApiClient(supertestUserWithCookieCredentials);
} else {
const kibanaServer = svlSharedConfig.get('servers.kibana');
const cAuthorities = svlSharedConfig.get('servers.kibana.certificateAuthorities');

const url = format({
...kibanaServer,
auth: false, // don't use auth in serverless
});
return createObservabilityAIAssistantApiClient(supertest.agent(url, { ca: cAuthorities }));
}
}

type ObservabilityAIAssistantApiClientKey = 'slsUser';
type ObservabilityAIAssistantApiClientKey = 'slsAdmin' | 'slsEditor' | 'slsUser';

export type ObservabilityAIAssistantApiClient = Record<
ObservabilityAIAssistantApiClientKey,
Awaited<ReturnType<typeof getObservabilityAIAssistantApiClient>>
>;
export function createObservabilityAIAssistantApiClient(st: supertest.Agent) {

export function createObservabilityAIAssistantApiClient(
st: SupertestWithRoleScope | supertest.Agent
) {
return <TEndpoint extends ObservabilityAIAssistantAPIEndpoint>(
options: {
type?: 'form-data';
endpoint: TEndpoint;
roleAuthc: RoleCredentials;
internalReqHeader: InternalRequestHeader;
roleAuthc?: RoleCredentials;
internalReqHeader?: InternalRequestHeader;
} & ObservabilityAIAssistantAPIClientRequestParamsOf<TEndpoint> & {
params?: { query?: { _inspect?: boolean } };
}
Expand All @@ -57,7 +67,8 @@ export function createObservabilityAIAssistantApiClient(st: supertest.Agent) {
const { method, pathname, version } = formatRequest(endpoint, params.path);
const url = format({ pathname, query: params?.query });

const headers: Record<string, string> = { ...internalReqHeader, ...roleAuthc.apiKeyHeader };
const headers: Record<string, string> =
roleAuthc && internalReqHeader ? { ...internalReqHeader, ...roleAuthc.apiKeyHeader } : {};

if (version) {
headers['Elastic-Api-Version'] = version;
Expand Down Expand Up @@ -182,10 +193,34 @@ export async function getObservabilityAIAssistantApiClientService({
getService,
}: InheritedFtrProviderContext): Promise<ObservabilityAIAssistantApiClient> {
const svlSharedConfig = getService('config');
// defaults to elastic_admin user when used without auth
const roleScopedSupertest = getService('roleScopedSupertest');

const supertestAdminWithCookieCredentials: SupertestWithRoleScope =
await roleScopedSupertest.getSupertestWithRoleScope('admin', {
useCookieHeader: true,
withInternalHeaders: true,
});

const supertestEditorWithCookieCredentials: SupertestWithRoleScope =
await roleScopedSupertest.getSupertestWithRoleScope('editor', {
useCookieHeader: true,
withInternalHeaders: true,
});

return {
// defaults to elastic_admin user when used without auth
slsUser: await getObservabilityAIAssistantApiClient({
svlSharedConfig,
}),
// cookie auth for internal apis
slsAdmin: await getObservabilityAIAssistantApiClient({
svlSharedConfig,
supertestUserWithCookieCredentials: supertestAdminWithCookieCredentials,
}),
// cookie auth for internal apis
slsEditor: await getObservabilityAIAssistantApiClient({
svlSharedConfig,
supertestUserWithCookieCredentials: supertestEditorWithCookieCredentials,
}),
};
}
47 changes: 23 additions & 24 deletions ...ss/api_integration/test_suites/observability/ai_assistant/tests/complete/complete.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import {
LlmResponseSimulator,
} from '@kbn/test-suites-xpack/observability_ai_assistant_api_integration/common/create_llm_proxy';
import { createOpenAiChunk } from '@kbn/test-suites-xpack/observability_ai_assistant_api_integration/common/create_openai_chunk';
import { SupertestWithRoleScope } from '@kbn/test-suites-xpack/api_integration/deployment_agnostic/services/role_scoped_supertest';
import { FtrProviderContext } from '../../common/ftr_provider_context';
import {
decodeEvents,
Expand All @@ -39,6 +40,9 @@ export default function ApiTest({ getService }: FtrProviderContext) {
const log = getService('log');
const svlUserManager = getService('svlUserManager');
const svlCommonApi = getService('svlCommonApi');
const roleScopedSupertest = getService('roleScopedSupertest');

let supertestEditorWithCookieCredentials: SupertestWithRoleScope;

const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantAPIClient');

Expand Down Expand Up @@ -82,10 +86,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
(body) => !isFunctionTitleRequest(body)
);
const responsePromise = new Promise<Response>((resolve, reject) => {
supertestWithoutAuth
supertestEditorWithCookieCredentials
.post(COMPLETE_API_URL)
.set(roleAuthc.apiKeyHeader)
.set(internalReqHeader)
.send({
messages,
connectorId,
Expand Down Expand Up @@ -134,6 +136,14 @@ export default function ApiTest({ getService }: FtrProviderContext) {
roleAuthc,
internalReqHeader,
});

supertestEditorWithCookieCredentials = await roleScopedSupertest.getSupertestWithRoleScope(
'editor',
{
useCookieHeader: true,
withInternalHeaders: true,
}
);
});

after(async () => {
Expand All @@ -155,10 +165,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {

const passThrough = new PassThrough();

supertestWithoutAuth
supertestEditorWithCookieCredentials
.post(COMPLETE_API_URL)
.set(roleAuthc.apiKeyHeader)
.set(internalReqHeader)
.send({
messages,
connectorId,
Expand Down Expand Up @@ -252,6 +260,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
},
});
});

describe('when creating a new conversation', () => {
let events: StreamingChatResponseEvent[];

Expand All @@ -270,12 +279,14 @@ export default function ApiTest({ getService }: FtrProviderContext) {
content: 'Hello',
},
});

expect(omit(events[1], 'id')).to.eql({
type: StreamingChatResponseEventType.ChatCompletionChunk,
message: {
content: ' again',
},
});

expect(omit(events[2], 'id', 'message.@timestamp')).to.eql({
type: StreamingChatResponseEventType.MessageAdd,
message: {
Expand Down Expand Up @@ -320,10 +331,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
)[0]?.conversation.id;

await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'DELETE /internal/observability_ai_assistant/conversation/{conversationId}',
roleAuthc,
internalReqHeader,
params: {
path: {
conversationId: createdConversationId,
Expand Down Expand Up @@ -398,10 +407,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
).to.eql(0);

const conversations = await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'POST /internal/observability_ai_assistant/conversations',
roleAuthc,
internalReqHeader,
})
.expect(200);

Expand Down Expand Up @@ -430,10 +437,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
.completeAfterIntercept();

const createResponse = await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'POST /internal/observability_ai_assistant/chat/complete',
roleAuthc,
internalReqHeader,
params: {
body: {
messages,
Expand All @@ -451,10 +456,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
conversationCreatedEvent = getConversationCreatedEvent(createResponse.body);

const conversationId = conversationCreatedEvent.conversation.id;
const fullConversation = await observabilityAIAssistantAPIClient.slsUser({
const fullConversation = await observabilityAIAssistantAPIClient.slsEditor({
endpoint: 'GET /internal/observability_ai_assistant/conversation/{conversationId}',
internalReqHeader,
roleAuthc,
params: {
path: {
conversationId,
Expand All @@ -467,10 +470,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
.completeAfterIntercept();

const updatedResponse = await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'POST /internal/observability_ai_assistant/chat/complete',
internalReqHeader,
roleAuthc,
params: {
body: {
messages: [
Expand Down Expand Up @@ -500,10 +501,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {

after(async () => {
await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'DELETE /internal/observability_ai_assistant/conversation/{conversationId}',
internalReqHeader,
roleAuthc,
params: {
path: {
conversationId: conversationCreatedEvent.conversation.id,
Expand Down
2 changes: 0 additions & 2 deletions ...ion/test_suites/observability/ai_assistant/tests/complete/functions/elasticsearch.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,6 @@ export default function ApiTest({ getService }: FtrProviderContext) {
const responseBody = await invokeChatCompleteWithFunctionRequest({
connectorId,
observabilityAIAssistantAPIClient,
internalReqHeader,
roleAuthc,
functionCall: {
name: ELASTICSEARCH_FUNCTION_NAME,
trigger: MessageRole.User,
Expand Down
9 changes: 1 addition & 8 deletions ...pi_integration/test_suites/observability/ai_assistant/tests/complete/functions/helpers.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ import {
} from '@kbn/observability-ai-assistant-plugin/common';
import type { AssistantScope } from '@kbn/ai-assistant-common';
import { Readable } from 'stream';
import type { InternalRequestHeader, RoleCredentials } from '../../../../../../../shared/services';
import { ObservabilityAIAssistantApiClient } from '../../../common/observability_ai_assistant_api_client';

function decodeEvents(body: Readable | string) {
Expand All @@ -34,22 +33,16 @@ export async function invokeChatCompleteWithFunctionRequest({
connectorId,
observabilityAIAssistantAPIClient,
functionCall,
roleAuthc,
internalReqHeader,
scopes,
}: {
connectorId: string;
observabilityAIAssistantAPIClient: ObservabilityAIAssistantApiClient;
functionCall: Message['message']['function_call'];
scopes?: AssistantScope[];
roleAuthc: RoleCredentials;
internalReqHeader: InternalRequestHeader;
}) {
const { body } = await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'POST /internal/observability_ai_assistant/chat/complete',
internalReqHeader,
roleAuthc,
params: {
body: {
messages: [
Expand Down
6 changes: 1 addition & 5 deletions ...gration/test_suites/observability/ai_assistant/tests/complete/functions/summarize.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,6 @@ export default function ApiTest({ getService }: FtrProviderContext) {
await invokeChatCompleteWithFunctionRequest({
connectorId,
observabilityAIAssistantAPIClient,
internalReqHeader,
roleAuthc,
functionCall: {
name: 'summarize',
trigger: MessageRole.User,
Expand All @@ -77,10 +75,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
});

it('persists entry in knowledge base', async () => {
const res = await observabilityAIAssistantAPIClient.slsUser({
const res = await observabilityAIAssistantAPIClient.slsEditor({
endpoint: 'GET /internal/observability_ai_assistant/kb/entries',
internalReqHeader,
roleAuthc,
params: {
query: {
query: '',
Expand Down
12 changes: 3 additions & 9 deletions ...pi_integration/test_suites/observability/ai_assistant/tests/connectors/connectors.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,19 +45,15 @@ export default function ApiTest({ getService }: FtrProviderContext) {

it('Returns a 2xx for enterprise license', async () => {
await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'GET /internal/observability_ai_assistant/connectors',
roleAuthc,
internalReqHeader,
})
.expect(200);
});

it('returns an empty list of connectors', async () => {
const res = await observabilityAIAssistantAPIClient.slsUser({
const res = await observabilityAIAssistantAPIClient.slsEditor({
endpoint: 'GET /internal/observability_ai_assistant/connectors',
roleAuthc,
internalReqHeader,
});

expect(res.body.length).to.be(0);
Expand All @@ -72,10 +68,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
roleAuthc,
});

const res = await observabilityAIAssistantAPIClient.slsUser({
const res = await observabilityAIAssistantAPIClient.slsEditor({
endpoint: 'GET /internal/observability_ai_assistant/connectors',
internalReqHeader,
roleAuthc,
});

expect(res.body.length).to.be(1);
Expand Down
Loading

0 comments on commit e4b0b80

Please sign in to comment.