Skip to content
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -81,18 +81,14 @@ describe('MockDataGeneratorModal', () => {
atlasAiService: {
getMockDataSchema: () => {
return Promise.resolve({
content: {
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.firstName',
fakerArgs: [],
isArray: false,
probability: 1.0,
},
],
},
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.firstName',
fakerArgs: [],
},
],
} as MockDataSchemaResponse);
},
},
Expand Down Expand Up @@ -287,42 +283,32 @@ describe('MockDataGeneratorModal', () => {
const mockServicesWithMockDataResponse = createMockServices();
mockServicesWithMockDataResponse.atlasAiService.getMockDataSchema = () =>
Promise.resolve({
content: {
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.firstName',
fakerArgs: [],
isArray: false,
probability: 1.0,
},
{
fieldPath: 'age',
mongoType: 'int',
fakerMethod: 'number.int',
fakerArgs: [],
isArray: false,
probability: 1.0,
},
{
fieldPath: 'email',
mongoType: 'string',
fakerMethod: 'internet',
fakerArgs: [],
isArray: false,
probability: 1.0,
},
{
fieldPath: 'username',
mongoType: 'string',
fakerMethod: 'noSuchMethod',
fakerArgs: [],
isArray: false,
probability: 1.0,
},
],
},
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.firstName',
fakerArgs: [],
},
{
fieldPath: 'age',
mongoType: 'int',
fakerMethod: 'number.int',
fakerArgs: [],
},
{
fieldPath: 'email',
mongoType: 'string',
fakerMethod: 'internet',
fakerArgs: [],
},
{
fieldPath: 'username',
mongoType: 'string',
fakerMethod: 'noSuchMethod',
fakerArgs: [],
},
],
});

it('shows a loading spinner when the faker schema generation is in progress', async () => {
Expand All @@ -332,9 +318,7 @@ describe('MockDataGeneratorModal', () => {
setTimeout(
() =>
resolve({
content: {
fields: [],
},
fields: [],
}),
1000
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,4 @@ export type MockDataGeneratorState =
| MockDataGeneratorCompletedState
| MockDataGeneratorErrorState;

export type FakerSchemaMapping = Omit<
MockDataSchemaResponse['content']['fields'][number],
'isArray'
>;
export type FakerSchemaMapping = MockDataSchemaResponse['fields'][number];
Original file line number Diff line number Diff line change
Expand Up @@ -702,7 +702,7 @@ const validateFakerSchema = (
fakerSchema: MockDataSchemaResponse,
logger: Logger
) => {
return fakerSchema.content.fields.map((field) => {
return fakerSchema.fields.map((field) => {
const { fakerMethod } = field;

const [moduleName, methodName, ...rest] = fakerMethod.split('.');
Expand Down
136 changes: 56 additions & 80 deletions packages/compass-generative-ai/src/atlas-ai-service.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -428,26 +428,20 @@ describe('AtlasAiService', function () {
if (apiURLPreset === 'cloud') {
it('makes a post request to the correct endpoint', async function () {
const mockResponse = {
content: {
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.fullName',
fakerArgs: [],
isArray: false,
probability: 1.0,
},
{
fieldPath: 'age',
mongoType: 'int',
fakerMethod: 'number.int',
fakerArgs: [{ json: '{"min": 18, "max": 65}' }],
isArray: false,
probability: 0.8,
},
],
},
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.fullName',
fakerArgs: [],
},
{
fieldPath: 'age',
mongoType: 'int',
fakerMethod: 'number.int',
fakerArgs: [{ json: '{"min": 18, "max": 65}' }],
},
],
};
const fetchStub = sandbox
.stub()
Expand All @@ -469,26 +463,20 @@ describe('AtlasAiService', function () {

it('includes sample values by default (includeSampleValues=true)', async function () {
const mockResponse = {
content: {
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.fullName',
fakerArgs: [],
isArray: false,
probability: 1.0,
},
{
fieldPath: 'age',
mongoType: 'int',
fakerMethod: 'number.int',
fakerArgs: [{ json: '{"min": 18, "max": 122}' }],
isArray: false,
probability: 0.8,
},
],
},
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.fullName',
fakerArgs: [],
},
{
fieldPath: 'age',
mongoType: 'int',
fakerMethod: 'number.int',
fakerArgs: [{ json: '{"min": 18, "max": 122}' }],
},
],
};
const fetchStub = sandbox
.stub()
Expand All @@ -515,26 +503,20 @@ describe('AtlasAiService', function () {

it('excludes sample values when includeSampleValues=false', async function () {
const mockResponse = {
content: {
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.fullName',
fakerArgs: [],
isArray: false,
probability: 1.0,
},
{
fieldPath: 'age',
mongoType: 'int',
fakerMethod: 'number.int',
fakerArgs: [{ json: '{"min": 18, "max": 65}' }],
isArray: false,
probability: 0.8,
},
],
},
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.fullName',
fakerArgs: [],
},
{
fieldPath: 'age',
mongoType: 'int',
fakerMethod: 'number.int',
fakerArgs: [{ json: '{"min": 18, "max": 65}' }],
},
],
};
const fetchStub = sandbox
.stub()
Expand All @@ -559,26 +541,20 @@ describe('AtlasAiService', function () {

it('makes POST request with correct headers and body structure', async function () {
const mockResponse = {
content: {
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.fullName',
fakerArgs: [],
isArray: false,
probability: 1.0,
},
{
fieldPath: 'age',
mongoType: 'int',
fakerMethod: 'number.int',
fakerArgs: [{ json: '{"min": 18, "max": 65}' }],
isArray: false,
probability: 0.8,
},
],
},
fields: [
{
fieldPath: 'name',
mongoType: 'string',
fakerMethod: 'person.fullName',
fakerArgs: [],
},
{
fieldPath: 'age',
mongoType: 'int',
fakerMethod: 'number.int',
fakerArgs: [{ json: '{"min": 18, "max": 65}' }],
},
],
};
const fetchStub = sandbox
.stub()
Expand Down
38 changes: 17 additions & 21 deletions packages/compass-generative-ai/src/atlas-ai-service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -232,27 +232,23 @@ export interface MockDataSchemaRequest {
}

export const MockDataSchemaResponseShape = z.object({
content: z.object({
fields: z.array(
z.object({
fieldPath: z.string(),
mongoType: z.string(),
fakerMethod: z.string(),
fakerArgs: z.array(
z.union([
z.object({
json: z.string(),
}),
z.string(),
z.number(),
z.boolean(),
])
),
isArray: z.boolean(),
probability: z.number(),
})
),
}),
fields: z.array(
z.object({
fieldPath: z.string(),
mongoType: z.string(),
fakerMethod: z.string(),
fakerArgs: z.array(
z.union([
z.object({
json: z.string(),
}),
z.string(),
z.number(),
z.boolean(),
])
),
})
),
});

export type MockDataSchemaResponse = z.infer<
Expand Down
Loading