diff --git a/src/Type/Php/PregMatchTypeSpecifyingExtension.php b/src/Type/Php/PregMatchTypeSpecifyingExtension.php
index 8678918..b46758b 100644
--- a/src/Type/Php/PregMatchTypeSpecifyingExtension.php
+++ b/src/Type/Php/PregMatchTypeSpecifyingExtension.php
@@ -1,5 +1,8 @@
 <?php declare(strict_types = 1);
 
+/*
+Blatantly copy-pasted from PHPStan's source code but with isFunctionSupported changed
+*/
 namespace TheCodingMachine\Safe\PHPStan\Type\Php;
 
 use PHPStan\Type\Php\RegexArrayShapeMatcher;
diff --git a/tests/Type/Php/ReplaceSafeFunctionsDynamicReturnTypeExtensionTest.php b/tests/Type/Php/ReplaceSafeFunctionsDynamicReturnTypeExtensionTest.php
deleted file mode 100644
index 3ebce33..0000000
--- a/tests/Type/Php/ReplaceSafeFunctionsDynamicReturnTypeExtensionTest.php
+++ /dev/null
@@ -1,22 +0,0 @@
-<?php
-
-namespace TheCodingMachine\Safe\PHPStan\Type\Php;
-
-use PHPUnit\Framework\TestCase;
-
-class ReplaceSafeFunctionsDynamicReturnTypeExtensionTest extends TestCase
-{
-    public function testWithStrings(): void
-    {
-        $x = \Safe\preg_replace('/foo/', 'bar', 'baz');
-
-        $this->assertStringNotContainsString('foo', $x);
-    }
-
-    public function testWithArrays(): void
-    {
-        $x = \Safe\preg_replace(['/foo/'], ['bar'], ['baz']);
-
-        $this->assertNotContains('foo', $x);
-    }
-}
diff --git a/tests/Type/Php/PregMatchParameterOutTypeExtensionTest.php b/tests/Type/Php/TypeAssertionsTest.php
similarity index 73%
rename from tests/Type/Php/PregMatchParameterOutTypeExtensionTest.php
rename to tests/Type/Php/TypeAssertionsTest.php
index 9a837b8..120b85a 100644
--- a/tests/Type/Php/PregMatchParameterOutTypeExtensionTest.php
+++ b/tests/Type/Php/TypeAssertionsTest.php
@@ -4,14 +4,16 @@
 
 use PHPStan\Testing\TypeInferenceTestCase;
 
-class PregMatchParameterOutTypeExtensionTest extends TypeInferenceTestCase
+class TypeAssertionsTest extends TypeInferenceTestCase
 {
     /**
      * @return iterable<mixed>
      */
     public static function dataFileAsserts(): iterable
     {
-        yield from self::gatherAssertTypes(__DIR__ . '/data/preg.php');
+        yield from self::gatherAssertTypes(__DIR__ . '/data/preg_match_unchecked.php');
+        yield from self::gatherAssertTypes(__DIR__ . '/data/preg_match_checked.php');
+        yield from self::gatherAssertTypes(__DIR__ . '/data/preg_replace_return.php');
     }
 
     /**
diff --git a/tests/Type/Php/data/preg.php b/tests/Type/Php/data/preg_match_checked.php
similarity index 61%
rename from tests/Type/Php/data/preg.php
rename to tests/Type/Php/data/preg_match_checked.php
index 40af91d..31699cb 100644
--- a/tests/Type/Php/data/preg.php
+++ b/tests/Type/Php/data/preg_match_checked.php
@@ -6,17 +6,6 @@
 $pattern = '/H(.)ll(o) (World)?/';
 $string = 'Hello World';
 
-// when return value isn't checked, we may-or-may-not have matches
-$type = "array{0?: string, 1?: non-empty-string, 2?: 'o', 3?: 'World'}";
-
-// @phpstan-ignore-next-line - use of unsafe is intentional
-\preg_match($pattern, $string, $matches);
-\PHPStan\Testing\assertType($type, $matches);
-
-\Safe\preg_match($pattern, $string, $matches);
-\PHPStan\Testing\assertType($type, $matches);
-
-
 // when the return value is checked, we should have matches,
 // unless the match-group itself is optional
 $type = "array{0: string, 1: non-empty-string, 2: 'o', 3?: 'World'}";
diff --git a/tests/Type/Php/data/preg_match_unchecked.php b/tests/Type/Php/data/preg_match_unchecked.php
new file mode 100644
index 0000000..8a55990
--- /dev/null
+++ b/tests/Type/Php/data/preg_match_unchecked.php
@@ -0,0 +1,17 @@
+<?php
+
+namespace TheCodingMachine\Safe\PHPStan\Type\Php\data;
+
+// Checking that preg_match and Safe\preg_match are equivalent
+$pattern = '/H(.)ll(o) (World)?/';
+$string = 'Hello World';
+
+// when return value isn't checked, we may-or-may-not have matches
+$type = "array{0?: string, 1?: non-empty-string, 2?: 'o', 3?: 'World'}";
+
+// @phpstan-ignore-next-line - use of unsafe is intentional
+\preg_match($pattern, $string, $matches);
+\PHPStan\Testing\assertType($type, $matches);
+
+\Safe\preg_match($pattern, $string, $matches);
+\PHPStan\Testing\assertType($type, $matches);
diff --git a/tests/Type/Php/data/preg_replace_return.php b/tests/Type/Php/data/preg_replace_return.php
new file mode 100644
index 0000000..94d6c8c
--- /dev/null
+++ b/tests/Type/Php/data/preg_replace_return.php
@@ -0,0 +1,11 @@
+<?php
+
+namespace TheCodingMachine\Safe\PHPStan\Type\Php\data;
+
+// preg_replace with a string pattern should return a string
+$x = \Safe\preg_replace('/foo/', 'bar', 'baz');
+\PHPStan\Testing\assertType("string", $x);
+
+// preg_replace with an array pattern should return an array
+$x = \Safe\preg_replace(['/foo/'], ['bar'], ['baz']);
+\PHPStan\Testing\assertType("array", $x);