diff --git a/entity_embed.module b/entity_embed.module
index ba1d1fd..4c76bc7 100644
--- a/entity_embed.module
+++ b/entity_embed.module
@@ -18,3 +18,54 @@ function entity_embed_help($route_name, RouteMatchInterface $route_match) {
       return '<p>' . t("Embed buttons define the buttons that can be added to CKEditor's toolbar. All buttons correspond to the selcted entity type, although a particular entity type can have more than one buttons. After you've created the desired buttons, go to <a href='@filterformats'>Text formats and editors configuration</a> page to add them to CKEditor's toolbar. Make sure to add <em>&lt;drupal-entity&gt;</em> to the list of 'Allowed HTML tags'.", array('@filterformats' => Url::fromRoute('filter.admin_overview')->toString())) . '</p>';
   }
 }
+
+/**
+ * Implements hook_embed_button_insert().
+ */
+function entity_embed_embed_button_insert(\Drupal\entity_embed\EmbedButtonInterface $entity) {
+  $button_icon_uuid = $entity->get('button_icon_uuid');
+  if ($button_icon_uuid) {
+    if ($file = \Drupal::entityManager()->loadEntityByUuid('file', $button_icon_uuid)) {
+      \Drupal::service('file.usage')->add($file, 'entity_embed', $entity->getEntityTypeId(), $entity->id());
+    }
+  }
+}
+
+/**
+ * Implements hook_embed_button_update().
+ */
+function entity_embed_embed_button_update(\Drupal\entity_embed\EmbedButtonInterface $entity) {
+  $new_button_icon_uuid = $entity->get('button_icon_uuid');
+
+  if (isset($entity->original)) {
+    $unchanged = $entity->original;
+    $old_button_icon_uuid = $unchanged->get('button_icon_uuid');
+
+    if (!empty($old_button_icon_uuid) && $old_button_icon_uuid != $new_button_icon_uuid) {
+      if ($file = \Drupal::entityManager()->loadEntityByUuid('file', $old_button_icon_uuid)) {
+        \Drupal::service('file.usage')->delete($file, 'entity_embed', $entity->getEntityTypeId(), $entity->id());
+      }
+    }
+  }
+
+  if ($new_button_icon_uuid) {
+    if ($file = \Drupal::entityManager()->loadEntityByUuid('file', $new_button_icon_uuid)) {
+      $usage = \Drupal::service('file.usage')->listUsage($file);
+      if (empty($usage['entity_embed'][$entity->getEntityTypeId()][$entity->id()])) {
+        \Drupal::service('file.usage')->add($file, 'entity_embed', $entity->getEntityTypeId(), $entity->id());
+      }
+    }
+  }
+}
+
+/**
+ * Implements hook_embed_button_update().
+ */
+function entity_embed_embed_button_delete(\Drupal\entity_embed\EmbedButtonInterface $entity) {
+  $button_icon_uuid = $entity->get('button_icon_uuid');
+  if ($button_icon_uuid) {
+    if ($file = \Drupal::entityManager()->loadEntityByUuid('file', $button_icon_uuid)) {
+      \Drupal::service('file.usage')->delete($file, 'entity_embed', $entity->getEntityTypeId(), $entity->id());
+    }
+  }
+}
diff --git a/src/Tests/EmbedButtonAdminTest.php b/src/Tests/EmbedButtonAdminTest.php
index 4dc8581..9d9e594 100644
--- a/src/Tests/EmbedButtonAdminTest.php
+++ b/src/Tests/EmbedButtonAdminTest.php
@@ -8,6 +8,7 @@
 namespace Drupal\entity_embed\Tests;
 
 Use \Drupal\Component\Utility\Unicode;
+use Drupal\file\Entity\File;
 
 /**
  * Tests the administrative UI.
@@ -99,4 +100,42 @@ class EmbedButtonAdminTest extends EntityEmbedTestBase {
     $this->assertNoText($name, 'Test embed_button does not appears on the list page');
   }
 
+  /**
+   * Tests the embed_button and file usage integration.
+   */
+  public function testEmbedButtonIcon() {
+    \Drupal::configFactory()->getEditable('system.file')->set('temporary_maximum_age', 1)->save();
+
+    $file1 = file_save_data(file_get_contents(drupal_get_path('module', 'entity_embed') . '/tests/fixture/entity.png'));
+    $file1->setTemporary();
+    $file1->save();
+
+    $file2 = file_save_data(file_get_contents(drupal_get_path('module', 'entity_embed') . '/tests/fixture/entity.png'));
+    $file2->setTemporary();
+    $file2->save();
+
+    $button = array(
+      'id' => 'test_button',
+      'label' => 'Testing embed button instance',
+      'button_label' => 'Test',
+      'entity_type' => 'node',
+      'entity_type_bundles' => array('article'),
+      'button_icon_uuid' => $file1->uuid(),
+      'display_plugins' => array('default'),
+    );
+
+    $entity = entity_create('embed_button', $button);
+    $entity->save();
+    $this->assertTrue(File::load($file1->id())->isPermanent());
+
+    $entity->button_icon_uuid = $file2->uuid();
+    $entity->save();
+
+    $this->assertTrue(File::load($file1->id())->isTemporary());
+    $this->assertTrue(File::load($file2->id())->isPermanent());
+
+    $entity->delete();
+    $this->assertTrue(File::load($file2->id())->isTemporary());
+  }
+
 }
diff --git a/tests/fixture/entity.png b/tests/fixture/entity.png
new file mode 100644
index 0000000..6960acd
Binary files /dev/null and b/tests/fixture/entity.png differ