From 300956719495b62cddf726e6d0fb7fbc322ae226 Mon Sep 17 00:00:00 2001
From: Jordan Welch <jordanhwelch@gmail.com>
Date: Sat, 26 Oct 2024 14:53:46 -0500
Subject: [PATCH] Add ability to extend `TestCall` to create custom chained
 helpers

---
 src/PendingCalls/TestCall.php     | 20 ++++++++++++++++++++
 tests/.snapshots/success.txt      |  6 +++++-
 tests/Features/TestCallExtend.php | 17 +++++++++++++++++
 tests/Visual/Parallel.php         |  2 +-
 4 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 tests/Features/TestCallExtend.php

diff --git a/src/PendingCalls/TestCall.php b/src/PendingCalls/TestCall.php
index a22ff11d..43aeac31 100644
--- a/src/PendingCalls/TestCall.php
+++ b/src/PendingCalls/TestCall.php
@@ -5,6 +5,7 @@
 namespace Pest\PendingCalls;
 
 use Closure;
+use Pest\Concerns\Extendable;
 use Pest\Concerns\Testable;
 use Pest\Exceptions\InvalidArgumentException;
 use Pest\Exceptions\TestDescriptionMissing;
@@ -31,6 +32,9 @@
 final class TestCall // @phpstan-ignore-line
 {
     use Describable;
+    use Extendable {
+        extend as traitExtend;
+    }
 
     /**
      * The list of test case factory attributes.
@@ -643,9 +647,25 @@ public function __get(string $name): self
      */
     public function __call(string $name, array $arguments): self
     {
+        if (self::hasExtend($name)) {
+            $extend = self::$extends[$name]->bindTo($this, TestCall::class);
+
+            if ($extend != false) { // @pest-arch-ignore-line
+                return $extend(...$arguments);
+            }
+        }
+
         return $this->addChain(Backtrace::file(), Backtrace::line(), $name, $arguments);
     }
 
+    public function extend(string $name, Closure $extend): void
+    {
+        $this->traitExtend($name, $extend);
+
+        $this->description = "extend: $name";
+        $this->testCaseMethod->closure = fn (): null => null;
+    }
+
     /**
      * Add a chain to the test case factory. Omitting the arguments will treat it as a property accessor.
      *
diff --git a/tests/.snapshots/success.txt b/tests/.snapshots/success.txt
index 1eaf4598..96906462 100644
--- a/tests/.snapshots/success.txt
+++ b/tests/.snapshots/success.txt
@@ -1324,6 +1324,10 @@
   ✓ a test
   ✓ higher order message test
 
+   PASS  Tests\Features\TestCallExtend
+  ✓ it uses fooBar extension with ('foo')
+  ✓ it uses fooBar extension with ('bar')
+
    PASS  Tests\Features\ThrowsNoExceptions
   ✓ it allows access to the underlying expectNotToPerformAssertions method
   ✓ it allows performing no expectations without being risky
@@ -1698,4 +1702,4 @@
    WARN  Tests\Visual\Version
   - visual snapshot of help command output
 
-  Tests:    2 deprecated, 4 warnings, 5 incomplete, 2 notices, 38 todos, 33 skipped, 1144 passed (2736 assertions)
\ No newline at end of file
+  Tests:    2 deprecated, 4 warnings, 5 incomplete, 2 notices, 38 todos, 33 skipped, 1146 passed (2738 assertions)
diff --git a/tests/Features/TestCallExtend.php b/tests/Features/TestCallExtend.php
new file mode 100644
index 00000000..91f4a6f5
--- /dev/null
+++ b/tests/Features/TestCallExtend.php
@@ -0,0 +1,17 @@
+<?php
+
+use function PHPUnit\Framework\assertTrue;
+
+test()->extend('fooBar', function () {
+    return $this->with([
+        'foo',
+        'bar',
+    ]);
+});
+
+it('uses fooBar extension', function ($value) {
+    assertTrue(in_array($value, [
+        'foo',
+        'bar',
+    ]));
+})->fooBar();
diff --git a/tests/Visual/Parallel.php b/tests/Visual/Parallel.php
index 313f8208..71cd6e83 100644
--- a/tests/Visual/Parallel.php
+++ b/tests/Visual/Parallel.php
@@ -16,7 +16,7 @@
 
 test('parallel', function () use ($run) {
     expect($run('--exclude-group=integration'))
-        ->toContain('Tests:    2 deprecated, 4 warnings, 5 incomplete, 2 notices, 38 todos, 24 skipped, 1134 passed (2712 assertions)')
+        ->toContain('Tests:    2 deprecated, 4 warnings, 5 incomplete, 2 notices, 38 todos, 24 skipped, 1136 passed (2714 assertions)')
         ->toContain('Parallel: 3 processes');
 })->skipOnWindows();