Create a cleaner way of managing course status

app/controllers/relationships_controller.rb

@@ -21,7 +21,7 @@ def create
       user.subscribe_to_section(section)
       render(json: { msg: "Subscribed" })
     else
-      user.follow(section)
+      user.follow_section(section)
       render(json: { msg: "Followed" })
     end
   end

app/models/relationship.rb

@@ -8,8 +8,49 @@ class Relationship < ApplicationRecord
   belongs_to :section
   has_one :review, dependent: :restrict_with_exception
 
+  enum stored_status: {
+    planned: "planned",
+    subscribed: "subscribed",
+    enrolled: "enrolled",
+  }
+
+  class Status < T::Enum
+    enums do
+      # A relationship can be in one of the following states:
+      Planned = new("planned") # A user has put this section into their class plan for a quarter
+      Subscribed = new("subscribed") # The section is in a class plan and the user wants enrollment notifications for the section
+      Enrolled = new("enrolled") # The section has been enrolled in. No need to send further notifications.
+
+      Completed = new("completed") # Where enrolled sections go after the quarter ends. The user has taken the course!
+      Reviewed = new("reviewed") # The user has taken the section and reviewed it.
+
+      # How does Status differ from stored_status above? We track the first three states (planned, subscribed, enrolled) in the database on this relationship model. The other two states (completed, reviewed) are derived based off other information (term calendar and if a review was written, respectively).
+    end
+  end
+
+  sig { returns(Status) }
+  def status
+    if reviewed?
+      Status::Reviewed
+    elsif completed?
+      Status::Completed
+    else
+      Status.deserialize(stored_status)
+    end
+  end
+
   sig { returns(T::Boolean) }
-  def review?
+  def reviewed?
     review.present?
   end
+
+  sig { returns(T::Boolean) }
+  def completed?
+    section&.term&.enrollment_period == Term::EnrollmentPeriod::Post
+  end
+
+  sig { returns(T::Boolean) }
+  def reviewed_or_completed?
+    reviewed? || completed?
+  end
 end

app/models/user.rb

@@ -113,23 +113,25 @@ def self.from_auth_payload(payload, referral_code:)
 
   # Marks a section as being followed by the given user.
   sig { params(section: Section).void }
-  def follow(section)
+  def follow_section(section)
     sections << section unless following?(section)
   end
 
   # Marks a section as being followed with notifications by the given user.
   sig { params(section: Section).void }
   def subscribe_to_section(section)
-    follow(section)
-    relationship = relationships.find_by(section:)
-    relationship&.update(notify: true) if use_notification_token
+    follow_section(section)
+    relationship = relationships.find_by!(section:)
+    raise ArgumentError, "Section cannot be subscribed to, it's past enrollment" if relationship.completed?
+
+    relationship.subscribed! if use_notification_token
   end
 
   # Whether a user is subscribed to a given section.
   sig { params(section: Section).returns(T::Boolean) }
   def subscribed_to_section?(section)
     relationship = relationships.find_by(section:)
-    relationship&.notify? || false
+    !!relationship&.subscribed?
   end
 
   # Whether a user is following a given section.
@@ -138,11 +140,17 @@ def following?(section)
     sections.include?(section)
   end
 
+  sig { params(section: Section).returns(Relationship::Status) }
+  def section_status(section)
+    relationship = relationships.find_by!(section:)
+    relationship.status
+  end
+
   # Stops following a section the user was following.
   sig { params(section: Section).void }
   def unfollow(section)
     relationship = relationships.find_by(section:)
-    raise ActiveRecord::RecordNotDestroyed, "Section is reviewed by user, cannot unfollow." if relationship&.review?
+    raise ActiveRecord::RecordNotDestroyed, "Section is reviewed by user, cannot unfollow." if relationship&.reviewed?
 
     sections.delete(section)
   end

db/data/20230827195318_move_notify_boolean_to_relationship_status.rb

@@ -0,0 +1,20 @@
+# typed: true
+# frozen_string_literal: true
+
+class MoveNotifyBooleanToRelationshipStatus < ActiveRecord::Migration[7.0]
+  def up
+    Relationship.all.find_each do |relationship|
+      # We have no way of telling whether a user enrolled in a section -- that's a new state!
+      relationship.stored_status = if relationship.notify
+                                     :subscribed
+                                   else
+                                     :planned
+                                   end
+      relationship.save!
+    end
+  end
+
+  def down
+    Relationship.update_all(stored_status: nil)
+  end
+end

db/data_schema.rb

@@ -0,0 +1 @@
+DataMigrate::Data.define(version: 20230827195318)

db/migrate/20230827154404_track_relationship_status.rb

@@ -0,0 +1,25 @@
+# typed: true
+# frozen_string_literal: true
+
+class TrackRelationshipStatus < ActiveRecord::Migration[7.0]
+  def change
+    reversible do |dir|
+      dir.up do
+        execute(<<-SQL.squish)
+          CREATE TYPE relationship_status AS ENUM ('planned', 'subscribed', 'enrolled');
+        SQL
+      end
+
+      dir.down do
+        execute(<<-SQL.squish)
+          DROP TYPE relationship_status;
+        SQL
+      end
+    end
+
+    change_table(:relationships, bulk: true) do |t|
+      # Will make this non-null in a future migration
+      t.column(:stored_status, :relationship_status, default: "planned")
+    end
+  end
+end

db/structure.sql

@@ -74,6 +74,17 @@ CREATE TYPE public.grade_type AS ENUM (
 );
 
 
+--
+-- Name: relationship_status; Type: TYPE; Schema: public; Owner: -
+--
+
+CREATE TYPE public.relationship_status AS ENUM (
+    'planned',
+    'subscribed',
+    'enrolled'
+);
+
+
 --
 -- Name: review_status; Type: TYPE; Schema: public; Owner: -
 --
@@ -521,6 +532,15 @@ CREATE TABLE public.courses_terms (
 );
 
 
+--
+-- Name: data_migrations; Type: TABLE; Schema: public; Owner: -
+--
+
+CREATE TABLE public.data_migrations (
+    version character varying NOT NULL
+);
+
+
 --
 -- Name: enrollment_appointments; Type: TABLE; Schema: public; Owner: -
 --
@@ -970,7 +990,8 @@ CREATE TABLE public.relationships (
     section_id bigint NOT NULL,
     created_at timestamp(6) without time zone NOT NULL,
     updated_at timestamp(6) without time zone NOT NULL,
-    notify boolean DEFAULT false NOT NULL
+    notify boolean DEFAULT false NOT NULL,
+    stored_status public.relationship_status DEFAULT 'planned'::public.relationship_status
 );
 
 
@@ -1633,6 +1654,14 @@ ALTER TABLE ONLY public.courses
     ADD CONSTRAINT courses_pkey PRIMARY KEY (id);
 
 
+--
+-- Name: data_migrations data_migrations_pkey; Type: CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public.data_migrations
+    ADD CONSTRAINT data_migrations_pkey PRIMARY KEY (version);
+
+
 --
 -- Name: enrollment_appointments enrollment_appointments_pkey; Type: CONSTRAINT; Schema: public; Owner: -
 --
@@ -2517,6 +2546,7 @@ INSERT INTO "schema_migrations" (version) VALUES
 ('20230430161016'),
 ('20230511033013'),
 ('20230718192047'),
-('20230820174337');
+('20230820174337'),
+('20230827154404');