From af26cd05d2dd57c49c673e4be1661710c10f65c8 Mon Sep 17 00:00:00 2001
From: Mark Skelton <mdskelton99@gmail.com>
Date: Thu, 19 Aug 2021 10:25:34 -0500
Subject: [PATCH] Support locators (#121)

* Add support for locators

* Fix locators

* Fix up types
---
 src/matchers/toBeChecked/index.test.ts | 12 ++++++++++++
 src/matchers/utils.ts                  | 26 ++++++++++++++++++--------
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/matchers/toBeChecked/index.test.ts b/src/matchers/toBeChecked/index.test.ts
index 7f102e0..6d23aa6 100644
--- a/src/matchers/toBeChecked/index.test.ts
+++ b/src/matchers/toBeChecked/index.test.ts
@@ -46,6 +46,18 @@ describe("toBeChecked", () => {
     })
   })
 
+  describe("locator", () => {
+    it("positive", async () => {
+      await page.setContent('<input type="radio" checked>')
+      await expect(page.locator("input")).toBeChecked()
+    })
+
+    it("positive", async () => {
+      await page.setContent('<input type="radio">')
+      await expect(page.locator("input")).not.toBeChecked()
+    })
+  })
+
   describe("with 'not' usage", () => {
     it("positive", async () => {
       await page.setContent('<input type="checkbox">')
diff --git a/src/matchers/utils.ts b/src/matchers/utils.ts
index c0c104c..549b757 100644
--- a/src/matchers/utils.ts
+++ b/src/matchers/utils.ts
@@ -1,16 +1,23 @@
-import type { Page, ElementHandle, Frame } from "playwright-core"
+import type { Page, ElementHandle, Frame, Locator } from "playwright-core"
 import { PageWaitForSelectorOptions } from "../../global"
 
-type Handle = Page | Frame | ElementHandle
+type Handle = Page | Frame | ElementHandle | Locator
 export type ExpectInputType = Handle | Promise<Handle>
 
 const isElementHandle = (value: Handle): value is ElementHandle => {
   return value.constructor.name === "ElementHandle"
 }
 
+const isLocator = (value: Handle): value is Locator => {
+  return value.constructor.name === "Locator"
+}
+
 export const getFrame = async (value: ExpectInputType) => {
   const resolved = await value
-  return isElementHandle(resolved) ? resolved.contentFrame() : resolved
+
+  return isElementHandle(resolved)
+    ? resolved.contentFrame()
+    : (resolved as Page | Frame)
 }
 
 const isObject = (value: unknown) =>
@@ -37,22 +44,25 @@ export const getElementHandle = async (
   const expectedValue = args.splice(-valueArgCount, valueArgCount) as string[]
 
   // Finally, we can find the element handle
-  const handle = await args[0]
-  let elementHandle = (await getFrame(handle)) ?? handle
+  let handle = await args[0]
+  handle = (await getFrame(handle)) ?? handle
 
+  if (isLocator(handle)) {
+    handle = (await handle.elementHandle())!
+  }
   // If the user provided a page or iframe, we need to locate the provided
   // selector or the `body` element if none was provided.
-  if (!isElementHandle(elementHandle)) {
+  else if (!isElementHandle(handle)) {
     const selector = args[1] ?? "body"
 
     try {
-      elementHandle = (await elementHandle.waitForSelector(selector, options))!
+      handle = (await handle.waitForSelector(selector, options))!
     } catch (err) {
       throw new Error(`Timeout exceed for element ${quote(selector)}`)
     }
   }
 
-  return [elementHandle, expectedValue] as const
+  return [handle, expectedValue] as const
 }
 
 export const quote = (val: string | null) => (val === null ? "" : `'${val}'`)