diff --git a/spec/system/assign_topic_spec.rb b/spec/system/assign_topic_spec.rb
index 455af25..4ece85c 100644
--- a/spec/system/assign_topic_spec.rb
+++ b/spec/system/assign_topic_spec.rb
@@ -30,6 +30,24 @@ describe "Assign | Assigning topics", type: :system, js: true do
       expect(page).to have_no_css("#topic .assigned-to")
     end
 
+    context "when assigns are not public" do
+      before { SiteSetting.assigns_public = false }
+
+      it "assigned small action post has 'private-assign' in class attribute" do
+        visit "/t/#{topic.id}"
+
+        topic_page.click_assign_topic
+        assign_modal.assignee = staff_user
+        assign_modal.confirm
+
+        expect(topic_page).to have_assigned(
+          user: staff_user,
+          at_post: 2,
+          class_attribute: ".private-assign",
+        )
+      end
+    end
+
     context "when unassign_on_close is set to true" do
       before { SiteSetting.unassign_on_close = true }
 
diff --git a/spec/system/page_objects/pages/topic.rb b/spec/system/page_objects/pages/topic.rb
index ab85b61..fe70463 100644
--- a/spec/system/page_objects/pages/topic.rb
+++ b/spec/system/page_objects/pages/topic.rb
@@ -27,7 +27,10 @@ module PageObjects
 
       def has_assignment_action?(args)
         assignee = args[:group]&.name || args[:user]&.username
-        container = args[:at_post] ? find("#post_#{args[:at_post]}") : page
+
+        container =
+          args[:at_post] ? find("#post_#{args[:at_post]}#{args[:class_attribute] || ""}") : page
+
         container.has_content?(
           I18n.t("js.action_codes.#{args[:action]}", who: "@#{assignee}", when: "just now"),
         )
diff --git a/test/javascripts/widgets/small-action-post-class-test.js b/test/javascripts/widgets/small-action-post-class-test.js
deleted file mode 100644
index 89d9eb4..0000000
--- a/test/javascripts/widgets/small-action-post-class-test.js
+++ /dev/null
@@ -1,58 +0,0 @@
-import { setupRenderingTest } from "discourse/tests/helpers/component-test";
-import { discourseModule, exists } from "discourse/tests/helpers/qunit-helpers";
-import hbs from "htmlbars-inline-precompile";
-import { render } from "@ember/test-helpers";
-import { test } from "qunit";
-import { withPluginApi } from "discourse/lib/plugin-api";
-import { resetPostSmallActionClassesCallbacks } from "discourse/widgets/post-small-action";
-
-discourseModule(
-  "Discourse Assign | Integration | Widget | Small Action Post Class",
-  function (hooks) {
-    setupRenderingTest(hooks);
-
-    test("Adds private-assign class when assigns are not public", async function (assert) {
-      try {
-        this.siteSettings.assigns_public = true;
-
-        this.set("args", {
-          id: 10,
-          actionCode: "assigned",
-        });
-
-        withPluginApi("1.6.0", (api) => {
-          api.addPostSmallActionClassesCallback((post) => {
-            if (
-              post.actionCode.includes("assigned") &&
-              !this.siteSettings.assigns_public
-            ) {
-              return ["private-assign"];
-            }
-          });
-        });
-
-        await render(
-          hbs`<MountWidget @widget="post-small-action" @args={{this.args}} />`
-        );
-
-        assert.notOk(
-          exists(".small-action.private-assign"),
-          "adds the private-assign class when assigns are public"
-        );
-
-        this.siteSettings.assigns_public = false;
-
-        await render(
-          hbs`<MountWidget @widget="post-small-action" @args={{this.args}} />`
-        );
-
-        assert.ok(
-          exists(".small-action.private-assign"),
-          "adds the private-assign class when assigns are NOT public"
-        );
-      } finally {
-        resetPostSmallActionClassesCallbacks();
-      }
-    });
-  }
-);