From 0edce4cf72bb3e3979c5e29665e81d11eab97033 Mon Sep 17 00:00:00 2001
From: Maryia Radchuk <maryia.radchuk@justeattakeaway.com>
Date: Fri, 22 Dec 2023 13:16:04 +0000
Subject: [PATCH] feat(pie-button): DSW-1539 wip assigning size prop to an icon
 slot

---
 packages/components/pie-button/src/index.ts | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/packages/components/pie-button/src/index.ts b/packages/components/pie-button/src/index.ts
index 41cff3bbf5..98c6287a4d 100644
--- a/packages/components/pie-button/src/index.ts
+++ b/packages/components/pie-button/src/index.ts
@@ -1,7 +1,7 @@
 import {
     LitElement, html, unsafeCSS, nothing, PropertyValues, TemplateResult,
 } from 'lit';
-import { property } from 'lit/decorators.js';
+import { property, queryAssignedElements } from 'lit/decorators.js';
 import { ifDefined } from 'lit/directives/if-defined.js';
 import { validPropertyValues, defineCustomElement } from '@justeattakeaway/pie-webc-core';
 import {
@@ -55,6 +55,22 @@ export class PieButton extends LitElement implements ButtonProps {
     updated (changedProperties: PropertyValues<this>): void {
         super.updated(changedProperties);
 
+        // Making sure that icon passed in a slot has the correct icon size
+        const [iconElement] = this._iconSlotElement;
+
+        if (iconElement) {
+            if (this.size === 'xsmall') {
+                iconElement.setAttribute('size', 'xs');
+            }
+            if (this.size === 'small-expressive' || this.size === 'small-productive') {
+                iconElement.setAttribute('size', 's');
+            }
+
+            if (this.size === 'medium' || this.size === 'large') {
+                iconElement.setAttribute('size', 'm');
+            }
+        }
+
         if (changedProperties.has('type')) {
             // If the new type is "submit", add the keydown event listener
             if (this.type === 'submit') {
@@ -117,6 +133,9 @@ export class PieButton extends LitElement implements ButtonProps {
     @property({ type: String })
     public responsiveSize?: ButtonProps['responsiveSize'];
 
+    @queryAssignedElements({ slot: 'icon' })
+        _iconSlotElement!: Array<HTMLElement>;
+
     /**
      * This method creates an invisible button of the same type as pie-button. It is then clicked, and immediately removed from the DOM.
      * This is done so that we trigger native form actions, such as submit and reset in the browser. The performance impact of adding and removing a single button to the DOM