fix: generate sequential ids for correlation rules

src/codegen/codegen.test.ts

@@ -13,6 +13,7 @@ import { checksRecording } from '@/test/fixtures/checksRecording'
 import { ThinkTime } from '@/types/testOptions'
 import { createProxyData, createRequest } from '@/test/factories/proxyData'
 import { jsonRule } from '@/test/fixtures/parameterizationRules'
+import { prettify } from '@/utils/prettify'
 
 const fakeDate = new Date('2000-01-01T00:00:00Z')
 
@@ -147,7 +148,7 @@ describe('Code generation', () => {
       ).toBe(expectedResult.replace(/\s/g, ''))
     })
 
-    it('should replace correlated values', () => {
+    it('should replace correlated values', async () => {
       const rules: TestRule[] = [
         {
           type: 'correlation',
@@ -161,34 +162,64 @@ describe('Code generation', () => {
             },
           },
         },
+        {
+          type: 'correlation',
+          id: '1',
+          extractor: {
+            filter: { path: '' },
+            selector: {
+              type: 'regex',
+              from: 'headers',
+              regex: 'project_id=(.*)$',
+            },
+          },
+        },
       ]
 
-      const expectedResult = `
-        params = { headers: {}, cookies: {} }
+      const expectedResult = await prettify(`
+        params = {
+          headers: {}, cookies: {}
+        }
+
         url = http.url\`http://test.k6.io/api/v1/foo\`
         resp = http.request('POST', url, null, params)
 
-        params = { headers: {}, cookies: {} }
-        url = http.url\`http://test.k6.io/api/v1/login\`
+        regex = new RegExp("project_id=(.*)$");
+        match = resp.headers["Project"].match(regex);
+        if (match) {
+          correlation_vars["correlation_0"] = match[1];
+        }
+
+        params = {
+          headers: {}, cookies: {}
+        }
+
+        url = http.url\`http://test.k6.io/api/v1/login?project_id=\${correlation_vars['correlation_0']}\`
         resp = http.request('POST', url, null, params)
-        correlation_vars['correlation_0'] = resp.json().user_id
 
-        params = { headers: {}, cookies: {} }
-        url = http.url\`http://test.k6.io/api/v1/users/\${correlation_vars['correlation_0']}\`
+        correlation_vars['correlation_1'] = resp.json().user_id
+
+        params = {
+          headers: {}, cookies: {}
+        }
+
+        url = http.url\`http://test.k6.io/api/v1/users/\${correlation_vars['correlation_1']}\`
         resp = http.request('GET', url, null, params)
 
-        params = { headers: {}, cookies: {} }
+        params = {
+          headers: {}, cookies: {}
+        }
+
         url = http.url\`http://test.k6.io/api/v1/users\`
-        resp = http.request('POST', url, \`${JSON.stringify({ user_id: "${correlation_vars['correlation_0']}" })}\`, params)
+        resp = http.request('POST', url, \`${JSON.stringify({ user_id: "${correlation_vars['correlation_1']}" })}\`, params)
 
-      `
+      `)
 
       expect(
-        generateRequestSnippets(correlationRecording, rules, thinkTime).replace(
-          /\s/g,
-          ''
+        await prettify(
+          generateRequestSnippets(correlationRecording, rules, thinkTime)
         )
-      ).toBe(expectedResult.replace(/\s/g, ''))
+      ).toBe(expectedResult)
     })
 
     it('should generate checks', () => {

src/rules/correlation.ts

@@ -19,97 +19,105 @@ import { replaceCorrelatedValues } from './correlation.utils'
 import { matchBeginEnd, matchRegex, getJsonObjectFromPath } from './shared'
 
 export function createCorrelationRuleInstance(
-  rule: CorrelationRule
+  rule: CorrelationRule,
+  idGenerator: Generator<number>
 ): CorrelationRuleInstance {
-  const sequentialIdGenerator = generateSequentialInt()
-
   const state: CorrelationState = {
     extractedValue: undefined,
     count: 0,
     responsesExtracted: [],
     requestsReplaced: [],
     generatedUniqueId: undefined,
-    sequentialIdGenerator,
+  }
+
+  function setState(newState: Partial<CorrelationState>) {
+    Object.assign(state, newState)
   }
 
   return {
     rule,
     state,
     type: rule.type,
-    apply: (
-      requestSnippetSchema: RequestSnippetSchema
-    ): RequestSnippetSchema => {
-      // this is the modified schema that we return to the accumulator
-      const snippetSchemaReturnValue = cloneDeep(requestSnippetSchema)
-
-      // if we have an extracted value we try to apply it to the request
-      // note: this comes before the extractor to avoid applying an extracted value on the same request/response pair
-      let uniqueId: number | undefined
-
-      if (state.extractedValue) {
-        // we populate uniqueId since it doesn't have to be regenerated
-        // this will be passed to the tryCorrelationExtraction function
-        uniqueId = state.generatedUniqueId
-
-        snippetSchemaReturnValue.data.request = replaceCorrelatedValues({
-          rule,
-          extractedValue: state.extractedValue,
-          uniqueId: uniqueId ?? 0,
-          request: requestSnippetSchema.data.request,
-        })
-
-        // Keep track of modified requests to display in preview
-        if (!isEqual(requestSnippetSchema, snippetSchemaReturnValue)) {
-          state.requestsReplaced = [
-            ...state.requestsReplaced,
-            {
-              original: requestSnippetSchema.data.request,
-              replaced: snippetSchemaReturnValue.data.request,
-            },
-          ]
-        }
-      }
-
-      // Skip extraction if filter doesn't match
-      if (!matchFilter(requestSnippetSchema, rule)) {
-        return snippetSchemaReturnValue
-      }
-
-      // try to extract the value
-      const {
-        extractedValue,
-        generatedUniqueId,
-        correlationExtractionSnippet,
-      } = tryCorrelationExtraction(
-        rule,
+    apply: (requestSnippetSchema: RequestSnippetSchema) =>
+      applyRule({ requestSnippetSchema, state, rule, idGenerator, setState }),
+  }
+}
+
+function applyRule({
+  requestSnippetSchema,
+  state,
+  rule,
+  idGenerator,
+  setState,
+}: {
+  requestSnippetSchema: RequestSnippetSchema
+  state: CorrelationState
+  rule: CorrelationRule
+  idGenerator: Generator<number>
+  setState: (newState: Partial<CorrelationState>) => void
+}) {
+  // this is the modified schema that we return to the accumulator
+  const snippetSchemaReturnValue = cloneDeep(requestSnippetSchema)
+
+  if (state.extractedValue) {
+    // we populate uniqueId since it doesn't have to be regenerated
+    // this will be passed to the tryCorrelationExtraction function
+
+    snippetSchemaReturnValue.data.request = replaceCorrelatedValues({
+      rule,
+      extractedValue: state.extractedValue,
+      uniqueId: state.generatedUniqueId ?? 0,
+      request: requestSnippetSchema.data.request,
+    })
+
+    // Keep track of modified requests to display in preview
+    if (!isEqual(requestSnippetSchema, snippetSchemaReturnValue)) {
+      setState({
+        requestsReplaced: [
+          ...state.requestsReplaced,
+          {
+            original: requestSnippetSchema.data.request,
+            replaced: snippetSchemaReturnValue.data.request,
+          },
+        ],
+      })
+    }
+  }
+
+  // Skip extraction if filter doesn't match
+  if (!matchFilter(requestSnippetSchema, rule)) {
+    return snippetSchemaReturnValue
+  }
+
+  // try to extract the value
+  const { extractedValue, generatedUniqueId, correlationExtractionSnippet } =
+    tryCorrelationExtraction(
+      rule,
+      requestSnippetSchema.data,
+      state.generatedUniqueId,
+      idGenerator
+    )
+
+  if (extractedValue && correlationExtractionSnippet) {
+    setState({
+      // Keep first extracted value
+      extractedValue: state.extractedValue ?? extractedValue,
+      generatedUniqueId: generatedUniqueId,
+      responsesExtracted: [
+        ...state.responsesExtracted,
         requestSnippetSchema.data,
-        uniqueId,
-        sequentialIdGenerator
-      )
+      ],
+
+      count: state.count + 1,
+    })
 
-      if (extractedValue && correlationExtractionSnippet) {
-        // Keep first extracted value
-
-        state.extractedValue ??= extractedValue
-        state.generatedUniqueId = generatedUniqueId
-        state.responsesExtracted = [
-          ...state.responsesExtracted,
-          requestSnippetSchema.data,
-        ]
-        state.count += 1
-
-        return {
-          ...snippetSchemaReturnValue,
-          after: [
-            ...requestSnippetSchema['after'],
-            correlationExtractionSnippet,
-          ],
-        }
-      }
-
-      return snippetSchemaReturnValue
-    },
+    return {
+      ...snippetSchemaReturnValue,
+      after: [...requestSnippetSchema['after'], correlationExtractionSnippet],
+    }
   }
+
+  return snippetSchemaReturnValue
 }
 
 const noCorrelationResult = {

src/rules/rules.ts

@@ -1,19 +1,17 @@
-import { RuleInstance, TestRule } from '@/types/rules'
+import { TestRule } from '@/types/rules'
 import { exhaustive } from '../utils/typescript'
 import { createCustomCodeRuleInstance } from './customCode'
 import { createCorrelationRuleInstance } from './correlation'
 import { createVerificationRuleInstance } from './verification'
 import { createParameterizationRuleInstance } from './parameterization'
 import { ProxyData, RequestSnippetSchema } from '@/types'
+import { generateSequentialInt } from './utils'
 
-export function applyRules(
-  recording: ProxyData[],
-  rules: TestRule[]
-): {
-  requestSnippetSchemas: RequestSnippetSchema[]
-  ruleInstances: RuleInstance[]
-} {
-  const ruleInstances = rules.map(createRuleInstance)
+export function applyRules(recording: ProxyData[], rules: TestRule[]) {
+  const idGenerator = generateSequentialInt()
+  const ruleInstances = rules.map((rule) =>
+    createRuleInstance(rule, idGenerator)
+  )
 
   const requestSnippetSchemas = recording.map((data) =>
     ruleInstances.reduce<RequestSnippetSchema>((acc, rule) => rule.apply(acc), {
@@ -26,10 +24,13 @@ export function applyRules(
   return { requestSnippetSchemas, ruleInstances }
 }
 
-export function createRuleInstance<T extends TestRule>(rule: T) {
+export function createRuleInstance<T extends TestRule>(
+  rule: T,
+  idGenerator: Generator<number>
+) {
   switch (rule.type) {
     case 'correlation':
-      return createCorrelationRuleInstance(rule)
+      return createCorrelationRuleInstance(rule, idGenerator)
     case 'parameterization':
       return createParameterizationRuleInstance(rule)
     case 'verification':

src/test/fixtures/correlationRecording.ts

@@ -23,7 +23,10 @@ export const correlationRecording: ProxyData[] = [
       path: '/api/v1/foo',
       reason: 'OK',
       httpVersion: '1.1',
-      headers: [['Content-Type', 'application/json']],
+      headers: [
+        ['Content-Type', 'application/json'],
+        ['project', 'project_id=555'],
+      ],
       cookies: [],
       content: JSON.stringify({ user_id: '444' }),
       contentLength: 0,
@@ -35,7 +38,7 @@ export const correlationRecording: ProxyData[] = [
     id: '2',
     request: {
       method: 'POST',
-      url: 'http://test.k6.io/api/v1/login',
+      url: 'http://test.k6.io/api/v1/login?project_id=555',
       headers: [],
       cookies: [],
       query: [],

src/types/rules.ts

@@ -32,7 +32,6 @@ export interface CorrelationState {
     replaced: Request
   }[]
   generatedUniqueId: number | undefined
-  sequentialIdGenerator: Generator<number>
 }
 
 export interface BaseRuleInstance<T extends TestRule> {

src/utils/prettify.ts

@@ -0,0 +1,11 @@
+import { format } from 'prettier/standalone'
+import * as prettierPluginBabel from 'prettier/plugins/babel'
+// eslint-disable-next-line import/namespace
+import * as prettierPluginEStree from 'prettier/plugins/estree'
+
+export function prettify(code: string) {
+  return format(code, {
+    parser: 'babel',
+    plugins: [prettierPluginBabel, prettierPluginEStree],
+  })
+}

src/views/Generator/Generator.utils.ts

@@ -1,8 +1,4 @@
 import { generateScript } from '@/codegen'
-import { format } from 'prettier/standalone'
-import * as prettierPluginBabel from 'prettier/plugins/babel'
-// eslint-disable-next-line import/namespace
-import * as prettierPluginEStree from 'prettier/plugins/estree'
 import { groupProxyData } from '@/utils/groups'
 import {
   selectFilteredRequests,
@@ -13,6 +9,7 @@ import { GeneratorFileData } from '@/types/generator'
 import { GeneratorFileDataSchema } from '@/schemas/generator'
 import { harToProxyData } from '@/utils/harToProxyData'
 import { GroupedProxyData } from '@/types'
+import { prettify } from '@/utils/prettify'
 
 export async function generateScriptPreview(
   generator: GeneratorFileData,
@@ -22,12 +19,8 @@ export async function generateScriptPreview(
     generator,
     recording,
   })
-  const prettifiedScript = await format(script, {
-    parser: 'babel',
-    plugins: [prettierPluginBabel, prettierPluginEStree],
-  })
 
-  return prettifiedScript
+  return prettify(script)
 }
 
 export function saveScript(script: string, fileName: string) {