diff --git a/frontend/test/token_tests/logic/token_actions_test.cljs b/frontend/test/token_tests/logic/token_actions_test.cljs
index 18711cb39e..e4679130a4 100644
--- a/frontend/test/token_tests/logic/token_actions_test.cljs
+++ b/frontend/test/token_tests/logic/token_actions_test.cljs
@@ -2,7 +2,9 @@
   (:require
    [app.common.test-helpers.compositions :as ctho]
    [app.common.test-helpers.files :as cthf]
-   [app.common.test-helpers.ids-map :as cthi]
+   [app.common.test-helpers.shapes :as cths]
+   [app.main.data.workspace.changes :as dch]
+   [app.main.data.workspace.selection :as dws]
    [app.main.ui.workspace.tokens.core :as wtc]
    [cljs.test :as t :include-macros true]
    [frontend-tests.helpers.pages :as thp]
@@ -11,45 +13,55 @@
 (t/use-fixtures :each
   {:before thp/reset-idmap!})
 
+
 (defn- setup-file
   []
-  (-> (cthf/sample-file :file-1)
-      (ctho/add-rect :file-1 :rect-1 {})
-      (ctho/add-rect :file-1 :rect-2 {})
-      (ctho/add-rect :file-1 :rect-3 {})))
+  (let [token-id (random-uuid)]
+    (-> (cthf/sample-file :file-1 :page-label :page-1)
+        (ctho/add-rect :rect-1 {})
+        (ctho/add-rect :rect-2 {})
+        (ctho/add-rect :rect-3 {})
+        (assoc-in [:data :tokens] {token-id {:id token-id
+                                             :value "12"
+                                             :name "sm"
+                                             :type :border-radius}}))))
 
 (t/deftest test-update-shape
   (t/async
     done
     (let [;; ==== Setup
-          file     (setup-file)
-          store    (ths/setup-store file)
+          file (setup-file)
+          store (ths/setup-store file)
+          rect-1 (cths/get-shape file :rect-1)
 
           ;; ==== Action
-          events [(wtc/on-add-token {:token-type-props {:attributes {:rx :ry}
-                                                        :on-update-shape #(fn [& _])}
-                                     :token {:id (random-uuid)}})]]
+          events [(dws/select-shape (:id rect-1))
+                  (dch/update-shapes [(:id rect-1)] (fn [shape] (assoc shape :applied-tokens {:rx (random-uuid)})))
+                  #_(wtc/on-add-token {:token-type-props {:attributes {:rx :ry}
+                                                          :on-update-shape #(fn [& _])}
+                                       :shape-ids [(:id rect-1)]
+                                       :token {:id (random-uuid)}})]]
       (ths/run-store
        store done events
        (fn [new-state]
-         (let [;; ==== Get
-               shape1' (get-in new-state [:workspace-data
-                                          :pages-index
-                                          (cthi/id :page1)
-                                          :objects
-                                          (cthi/id :shape1)])
-               fills'      (:fills shape1')
-               fill'       (first fills')]
+         (let [file' (ths/get-file-from-store new-state)
+               page' (cthf/current-page file')
+               rect-1' (cths/get-shape file' :rect-1)
+               ;; ==== Get
+               #_#_rect-1' (get-in new-state [:workspace-data
+                                              :pages-index
+                                              (cthi/id :page-1)
+                                              :objects
+                                              (cthi/id :rect-1)])]
+
+           ;; ==== Check
+           (t/is (some? (:applied-tokens rect-1')))))))))
 
-          ;; ==== Check
-           (t/is (some? shape1'))
-           (t/is (= (count fills') 1))
-           (t/is (= (:fill-color fill') "#fabada"))
-           (t/is (= (:fill-opacity fill') 1))))))))
 
 
 
 (comment
+  (t/run-tests)
   (defn make-printable
     "Convert records that are not printable by cider inspect into regular maps."
     [coll]