From 3cccc52515d03f2dae462fbbf3db5255c6e7fc47 Mon Sep 17 00:00:00 2001
From: Jan Jakes <jan.jakes@automattic.com>
Date: Tue, 5 Jul 2022 14:59:26 +0200
Subject: [PATCH] Require subjects by class name to have them type checked and
 autocompleted

[MAILPOET-4465]
---
 mailpoet/lib/Automation/Engine/Exceptions.php |  6 ++++
 .../Engine/Workflows/WorkflowRun.php          | 31 +++++++++++++++++--
 .../Actions/SendWelcomeEmailAction.php        |  4 +--
 3 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/mailpoet/lib/Automation/Engine/Exceptions.php b/mailpoet/lib/Automation/Engine/Exceptions.php
index 19b3f1ede3..e904757cce 100644
--- a/mailpoet/lib/Automation/Engine/Exceptions.php
+++ b/mailpoet/lib/Automation/Engine/Exceptions.php
@@ -95,6 +95,12 @@ class Exceptions {
       ->withMessage(__(sprintf("Subject with key '%s' not found.", $key), 'mailpoet'));
   }
 
+  public static function subjectClassNotFound(string $key): NotFoundException {
+    return NotFoundException::create()
+      ->withErrorCode(self::SUBJECT_NOT_FOUND)
+      ->withMessage(__(sprintf("Subject of class '%s' not found.", $key), 'mailpoet'));
+  }
+
   public static function subjectLoadFailed(string $key, array $args): InvalidStateException {
     return InvalidStateException::create()
       ->withErrorCode(self::SUBJECT_LOAD_FAILED)
diff --git a/mailpoet/lib/Automation/Engine/Workflows/WorkflowRun.php b/mailpoet/lib/Automation/Engine/Workflows/WorkflowRun.php
index f421cf4f09..71fbaaa800 100644
--- a/mailpoet/lib/Automation/Engine/Workflows/WorkflowRun.php
+++ b/mailpoet/lib/Automation/Engine/Workflows/WorkflowRun.php
@@ -4,6 +4,7 @@ namespace MailPoet\Automation\Engine\Workflows;
 
 use DateTimeImmutable;
 use MailPoet\Automation\Engine\Exceptions;
+use MailPoet\Automation\Engine\Exceptions\InvalidStateException;
 use MailPoet\Automation\Engine\Utils\Json;
 
 class WorkflowRun {
@@ -33,6 +34,12 @@ class WorkflowRun {
   /** @var Subject[] */
   private $subjects;
 
+  /** @var array<class-string, string> */
+  private $subjectKeyClassMap = [];
+
+  /**
+   * @param Subject[] $subjects
+   */
   public function __construct(
     int $workflowId,
     string $triggerKey,
@@ -43,6 +50,10 @@ class WorkflowRun {
     $this->triggerKey = $triggerKey;
     $this->subjects = $subjects;
 
+    foreach ($subjects as $subject) {
+      $this->subjectKeyClassMap[get_class($subject)] = $subject->getKey();
+    }
+
     if ($id) {
       $this->id = $id;
     }
@@ -88,7 +99,17 @@ class WorkflowRun {
     return $this->subjects;
   }
 
-  public function requireSubject(string $key): Subject {
+  /**
+   * @template T of Subject
+   * @param class-string<T> $class
+   * @return T
+   */
+  public function requireSubject(string $class): Subject {
+    $key = $this->subjectKeyClassMap[$class] ?? null;
+    if (!$key) {
+      throw Exceptions::subjectClassNotFound($class);
+    }
+
     $subjects = $this->getSubjects($key);
     if (count($subjects) === 0) {
       throw Exceptions::subjectNotFound($key);
@@ -96,7 +117,13 @@ class WorkflowRun {
     if (count($subjects) > 1) {
       throw Exceptions::multipleSubjectsFound($key);
     }
-    return $subjects[0];
+
+    // ensure PHPStan we're indeed returning an instance of $class
+    $subject = $subjects[0];
+    if (!$subject instanceof $class) {
+      throw new InvalidStateException();
+    }
+    return $subject;
   }
 
   public function toArray(): array {
diff --git a/mailpoet/lib/Automation/Integrations/MailPoet/Actions/SendWelcomeEmailAction.php b/mailpoet/lib/Automation/Integrations/MailPoet/Actions/SendWelcomeEmailAction.php
index 4801bc39e6..da6eebd099 100644
--- a/mailpoet/lib/Automation/Integrations/MailPoet/Actions/SendWelcomeEmailAction.php
+++ b/mailpoet/lib/Automation/Integrations/MailPoet/Actions/SendWelcomeEmailAction.php
@@ -76,7 +76,7 @@ class SendWelcomeEmailAction implements Action {
   public function run(Workflow $workflow, WorkflowRun $workflowRun, Step $step): void {
     $newsletter = $this->getWelcomeEmailForStep($step);
 
-    $subscriberSubject = $workflowRun->requireSubject(SubscriberSubject::KEY);
+    $subscriberSubject = $workflowRun->requireSubject(SubscriberSubject::class);
     $subscriberId = $subscriberSubject->getFields()['id']->getValue();
     $subscriber = $this->subscribersRepository->findOneById($subscriberId);
 
@@ -88,7 +88,7 @@ class SendWelcomeEmailAction implements Action {
       throw InvalidStateException::create()->withMessage(sprintf("Cannot schedule a newsletter for subscriber ID '%s' because their status is '%s'.", $subscriber->getId(), $subscriber->getStatus()));
     }
 
-    $segmentSubject = $workflowRun->requireSubject(SegmentSubject::KEY);
+    $segmentSubject = $workflowRun->requireSubject(SegmentSubject::class);
     $segmentId = $segmentSubject->getFields()['id']->getValue();
     $subscriberSegment = $this->subscriberSegmentRepository->findOneBy([
       'subscriber' => $subscriber,