From 9e0aa67d286b106600cae9c8837d5e4797b9a8f7 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Wed, 4 May 2022 11:14:53 +0100
Subject: [PATCH] feat(codegen): brush up context options in pytest codegen
 (#13924)

---
 .../src/server/recorder/python.ts             | 46 ++++++++++---------
 .../inspector/cli-codegen-pytest.spec.ts      | 12 ++---
 2 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/packages/playwright-core/src/server/recorder/python.ts b/packages/playwright-core/src/server/recorder/python.ts
index 69dc0a6be8..c7c90357de 100644
--- a/packages/playwright-core/src/server/recorder/python.ts
+++ b/packages/playwright-core/src/server/recorder/python.ts
@@ -33,20 +33,20 @@ export class PythonLanguageGenerator implements LanguageGenerator {
   private _awaitPrefix: '' | 'await ';
   private _asyncPrefix: '' | 'async ';
   private _isAsync: boolean;
-  private _isTest: boolean;
+  private _isPyTest: boolean;
 
-  constructor(isAsync: boolean, isTest: boolean) {
-    this.id = isTest ? 'pytest' : (isAsync ? 'python-async' : 'python');
-    this.fileName = isTest ? 'Pytest' : (isAsync ? 'Python Async' : 'Python');
+  constructor(isAsync: boolean, isPyTest: boolean) {
+    this.id = isPyTest ? 'pytest' : (isAsync ? 'python-async' : 'python');
+    this.fileName = isPyTest ? 'Pytest' : (isAsync ? 'Python Async' : 'Python');
     this._isAsync = isAsync;
-    this._isTest = isTest;
+    this._isPyTest = isPyTest;
     this._awaitPrefix = isAsync ? 'await ' : '';
     this._asyncPrefix = isAsync ? 'async ' : '';
   }
 
   generateAction(actionInContext: ActionInContext): string {
     const action = actionInContext.action;
-    if (this._isTest && (action.name === 'openPage' || action.name === 'closePage'))
+    if (this._isPyTest && (action.name === 'openPage' || action.name === 'closePage'))
       return '';
 
     const pageAlias = actionInContext.frame.pageAlias;
@@ -155,20 +155,18 @@ export class PythonLanguageGenerator implements LanguageGenerator {
 
   generateHeader(options: LanguageGeneratorOptions): string {
     const formatter = new PythonFormatter();
-    if (this._isTest) {
-      formatter.add(`${options.deviceName ? 'import pytest\n' : ''}
-from playwright.sync_api import Page, expect
-${options.deviceName ? `
+    if (this._isPyTest) {
+      const contextOptions = formatContextOptions(options.contextOptions, options.deviceName, true /* asDict */);
+      const fixture = contextOptions ? `
 
 @pytest.fixture(scope="session")
 def browser_context_args(browser_context_args, playwright) {
-    device = playwright.devices["${options.deviceName}"]
-    return dict(
-        **browser_context_args,
-        **device,
-    )
+    return {${contextOptions}}
 }
-` : ''}
+` : '';
+      formatter.add(`${options.deviceName ? 'import pytest\n' : ''}
+from playwright.sync_api import Page, expect
+${fixture}
 
 def test_example(page: Page) -> None {`);
     } else if (this._isAsync) {
@@ -194,7 +192,7 @@ def run(playwright: Playwright) -> None {
   }
 
   generateFooter(saveStorage: string | undefined): string {
-    if (this._isTest) {
+    if (this._isPyTest) {
       return '';
     } else if (this._isAsync) {
       const storageStateLine = saveStorage ? `\n    await context.storage_state(path=${quote(saveStorage)})` : '';
@@ -245,18 +243,22 @@ function toSnakeCase(name: string): string {
   return name.replace(toSnakeCaseRegex, `_$1`).toLowerCase();
 }
 
-function formatOptions(value: any, hasArguments: boolean): string {
+function formatOptions(value: any, hasArguments: boolean, asDict?: boolean): string {
   const keys = Object.keys(value);
   if (!keys.length)
     return '';
-  return (hasArguments ? ', ' : '') + keys.map(key => `${toSnakeCase(key)}=${formatValue(value[key])}`).join(', ');
+  return (hasArguments ? ', ' : '') + keys.map(key => {
+    if (asDict)
+      return `"${toSnakeCase(key)}": ${formatValue(value[key])}`;
+    return `${toSnakeCase(key)}=${formatValue(value[key])}`;
+  }).join(', ');
 }
 
-function formatContextOptions(options: BrowserContextOptions, deviceName: string | undefined): string {
+function formatContextOptions(options: BrowserContextOptions, deviceName: string | undefined, asDict?: boolean): string {
   const device = deviceName && deviceDescriptors[deviceName];
   if (!device)
-    return formatOptions(options, false);
-  return `**playwright.devices[${quote(deviceName!)}]` + formatOptions(sanitizeDeviceOptions(device, options), true);
+    return formatOptions(options, false, asDict);
+  return `**playwright.devices[${quote(deviceName!)}]` + formatOptions(sanitizeDeviceOptions(device, options), true, asDict);
 }
 
 class PythonFormatter {
diff --git a/tests/library/inspector/cli-codegen-pytest.spec.ts b/tests/library/inspector/cli-codegen-pytest.spec.ts
index d3daf64e5c..03728cbcde 100644
--- a/tests/library/inspector/cli-codegen-pytest.spec.ts
+++ b/tests/library/inspector/cli-codegen-pytest.spec.ts
@@ -30,11 +30,11 @@ def test_example(page: Page) -> None:`;
   expect(cli.text()).toContain(expectedResult);
 });
 
-test('should print the correct context options when using a device', async ({ browserName, runCLI }, testInfo) => {
+test('should print the correct context options when using a device and lang', async ({ browserName, runCLI }, testInfo) => {
   test.skip(browserName !== 'webkit');
 
   const tmpFile = testInfo.outputPath('script.js');
-  const cli = runCLI(['--target=pytest', '--device=iPhone 11', '--output', tmpFile, emptyHTML]);
+  const cli = runCLI(['--target=pytest', '--device=iPhone 11', '--lang=en-US', '--output', tmpFile, emptyHTML]);
   await cli.exited;
   const content = fs.readFileSync(tmpFile);
   expect(content.toString()).toBe(`import pytest
@@ -44,11 +44,7 @@ from playwright.sync_api import Page, expect
 
 @pytest.fixture(scope="session")
 def browser_context_args(browser_context_args, playwright):
-    device = playwright.devices["iPhone 11"]
-    return dict(
-    **browser_context_args,
-    **device,
-    )
+    return {**playwright.devices["iPhone 11"], "locale": "en-US"}
 
 
 def test_example(page: Page) -> None:
@@ -59,7 +55,7 @@ def test_example(page: Page) -> None:
 });
 
 test('should save the codegen output to a file if specified', async ({ runCLI }, testInfo) => {
-  const tmpFile = testInfo.outputPath('script.js');
+  const tmpFile = testInfo.outputPath('test_example.py');
   const cli = runCLI(['--target=pytest', '--output', tmpFile, emptyHTML]);
   await cli.exited;
   const content = fs.readFileSync(tmpFile);