Encapsulate enrichment actions, add tests

app/models/concerns/enrichable.rb

@@ -22,16 +22,23 @@ module Enrichable
     }
   end
 
-  def log_enrichment!(attribute_name:, attribute_value:, source:, metadata: {})
-    de = DataEnrichment.find_or_create_by!(
-      enrichable: self,
-      attribute_name: attribute_name,
-      source: source,
-    )
+  # Convenience method for a single attribute
+  def enrich_attribute(attr, value, source:, metadata: {})
+    enrich_attributes({ attr => value }, source:, metadata:)
+  end
 
-    de.value = attribute_value
-    de.metadata = metadata
-    de.save!
+  # Enriches all attributes that haven't been locked yet
+  def enrich_attributes(attrs, source:, metadata: {})
+    enrichable_attrs = Array(attrs).reject { |k, _v| locked?(k) }
+
+    ActiveRecord::Base.transaction do
+      enrichable_attrs.each do |attr, value|
+        self.send("#{attr}=", value)
+        log_enrichment(attribute_name: attr, attribute_value: value, source: source, metadata: metadata)
+      end
+
+      save
+    end
   end
 
   def locked?(attr)
@@ -57,6 +64,18 @@ module Enrichable
   end
 
   private
+    def log_enrichment(attribute_name:, attribute_value:, source:, metadata: {})
+      de = DataEnrichment.find_or_create_by(
+        enrichable: self,
+        attribute_name: attribute_name,
+        source: source,
+      )
+
+      de.value = attribute_value
+      de.metadata = metadata
+      de.save
+    end
+
     def ignored_enrichable_attributes
       %w[id updated_at created_at]
     end

app/models/family/auto_categorizer.rb

@@ -27,23 +27,19 @@ class Family::AutoCategorizer
     end
 
     scope.each do |transaction|
-      transaction.lock_attr!(:category_id)
-
       auto_categorization = result.data.find { |c| c.transaction_id == transaction.id }
 
       category_id = user_categories_input.find { |c| c[:name] == auto_categorization&.category_name }&.dig(:id)
 
       if category_id.present?
-        Family.transaction do
-          transaction.log_enrichment!(
-            attribute_name: "category_id",
-            attribute_value: category_id,
-            source: "ai",
-          )
-
-          transaction.update!(category_id: category_id)
-        end
+        transaction.enrich_attribute(
+          :category_id,
+          category_id,
+          source: "ai"
+        )
       end
+
+      transaction.lock_attr!(:category_id)
     end
   end
 

app/models/family/auto_merchant_detector.rb

@@ -27,8 +27,6 @@ class Family::AutoMerchantDetector
     end
 
     scope.each do |transaction|
-      transaction.lock_attr!(:merchant_id)
-
       auto_detection = result.data.find { |c| c.transaction_id == transaction.id }
 
       merchant_id = user_merchants_input.find { |m| m[:name] == auto_detection&.business_name }&.dig(:id)
@@ -46,16 +44,16 @@ class Family::AutoMerchantDetector
       merchant_id = merchant_id || ai_provider_merchant&.id
 
       if merchant_id.present?
-        Family.transaction do
-          transaction.log_enrichment!(
-            attribute_name: "merchant_id",
-            attribute_value: merchant_id,
-            source: "ai",
-          )
+        transaction.enrich_attribute(
+          :merchant_id,
+          merchant_id,
+          source: "ai"
+        )
 
-          transaction.update!(merchant_id: merchant_id)
-        end
       end
+
+      # We lock the attribute so that this Rule doesn't try to run again
+      transaction.lock_attr!(:merchant_id)
     end
   end
 

app/models/plaid_item.rb

@@ -77,14 +77,14 @@ class PlaidItem < ApplicationRecord
         category = alias_matcher.match(transaction.plaid_category_detailed)
 
         if category.present?
-          PlaidItem.transaction do
-            transaction.log_enrichment!(
-              attribute_name: "category_id",
-              attribute_value: category.id,
-              source: "plaid"
-            )
-            transaction.set_category!(category)
+          # Matcher could either return a string or a Category object
+          user_category = if category.is_a?(String)
+            family.categories.find_or_create_by!(name: category)
+          else
+            category
           end
+
+          transaction.enrich_attribute(:category_id, user_category.id, source: "plaid")
         end
       end
     end

app/models/rule/action_executor/set_transaction_category.rb

@@ -17,15 +17,11 @@ class Rule::ActionExecutor::SetTransactionCategory < Rule::ActionExecutor
     end
 
     scope.each do |txn|
-      Rule.transaction do
-        txn.log_enrichment!(
-          attribute_name: "category_id",
-          attribute_value: category.id,
-          source: "rule"
-        )
-
-        txn.update!(category: category)
-      end
+      txn.enrich_attribute(
+        :category_id,
+        category.id,
+        source: "rule"
+      )
     end
   end
 end

app/models/rule/action_executor/set_transaction_merchant.rb

@@ -17,14 +17,11 @@ class Rule::ActionExecutor::SetTransactionMerchant < Rule::ActionExecutor
     end
 
     scope.each do |txn|
-      Rule.transaction do
-        txn.log_enrichment!(
-          attribute_name: "merchant_id",
-          attribute_value: merchant.id,
-          source: "rule"
-        )
-        txn.update!(merchant: merchant)
-      end
+      txn.enrich_attribute(
+        :merchant_id,
+        merchant.id,
+        source: "rule"
+      )
     end
   end
 end

app/models/rule/action_executor/set_transaction_name.rb

@@ -16,14 +16,11 @@ class Rule::ActionExecutor::SetTransactionName < Rule::ActionExecutor
     end
 
     scope.each do |txn|
-      Rule.transaction do
-        txn.entry.log_enrichment!(
-          attribute_name: "name",
-          attribute_value: value,
-          source: "rule"
-        )
-        txn.entry.update!(name: value)
-      end
+      txn.entry.enrich_attribute(
+        :name,
+        value,
+        source: "rule"
+      )
     end
   end
 end

app/models/rule/action_executor/set_transaction_tags.rb

@@ -17,15 +17,11 @@ class Rule::ActionExecutor::SetTransactionTags < Rule::ActionExecutor
     end
 
     rows = scope.each do |txn|
-      Rule.transaction do
-        txn.log_enrichment!(
-          attribute_name: "tag_ids",
-          attribute_value: [ tag.id ],
-          source: "rule"
-        )
-
-        txn.update!(tag_ids: [ tag.id ])
-      end
+      txn.enrich_attribute(
+        :tag_ids,
+        [ tag.id ],
+        source: "rule"
+      )
     end
   end
 end

test/models/concerns/enrichable_test.rb

@@ -0,0 +1,79 @@
+require "test_helper"
+
+class EnrichableTest < ActiveSupport::TestCase
+  setup do
+    @enrichable = accounts(:depository)
+  end
+
+  test "can enrich multiple attributes" do
+    assert_difference "DataEnrichment.count", 2 do
+      @enrichable.enrich_attributes({ name: "Updated Checking", balance: 6_000 }, source: "plaid")
+    end
+
+    assert_equal "Updated Checking", @enrichable.name
+    assert_equal 6_000, @enrichable.balance.to_d
+  end
+
+  test "can enrich a single attribute" do
+    assert_difference "DataEnrichment.count", 1 do
+      @enrichable.enrich_attribute(:name, "Single Update", source: "ai")
+    end
+
+    assert_equal "Single Update", @enrichable.name
+  end
+
+  test "can lock an attribute" do
+    refute @enrichable.locked?(:name)
+
+    @enrichable.lock_attr!(:name)
+    assert @enrichable.locked?(:name)
+  end
+
+  test "can unlock an attribute" do
+    @enrichable.lock_attr!(:name)
+    assert @enrichable.locked?(:name)
+
+    @enrichable.unlock_attr!(:name)
+    refute @enrichable.locked?(:name)
+  end
+
+  test "can lock saved attributes" do
+    @enrichable.name = "User Override"
+    @enrichable.balance = 1_234
+    @enrichable.save!
+
+    @enrichable.lock_saved_attributes!
+
+    assert @enrichable.locked?(:name)
+    assert @enrichable.locked?(:balance)
+  end
+
+  test "does not enrich locked attributes" do
+    original_name = @enrichable.name
+
+    @enrichable.lock_attr!(:name)
+
+    assert_no_difference "DataEnrichment.count" do
+      @enrichable.enrich_attribute(:name, "Should Not Change", source: "plaid")
+    end
+
+    assert_equal original_name, @enrichable.reload.name
+  end
+
+  test "enrichable? reflects lock state" do
+    assert @enrichable.enrichable?(:name)
+
+    @enrichable.lock_attr!(:name)
+
+    refute @enrichable.enrichable?(:name)
+  end
+
+  test "enrichable scope includes and excludes records based on lock state" do
+    # Initially, the record should be enrichable for :name
+    assert_includes Account.enrichable(:name), @enrichable
+
+    @enrichable.lock_attr!(:name)
+
+    refute_includes Account.enrichable(:name), @enrichable
+  end
+end