From 8d5a1cd774c3d3147b4464e6676092c7a61779e3 Mon Sep 17 00:00:00 2001
From: Shayne Chen <shayne.chen@mediatek.com>
Date: Fri, 25 Feb 2022 09:36:01 +0800
Subject: [PATCH] mt76: testmode: rework tx antenna setting

---
 mt7915/mcu.c      | 7 +------
 mt7915/testmode.c | 9 +--------
 testmode.c        | 4 ++--
 3 files changed, 4 insertions(+), 16 deletions(-)

diff --git a/mt7915/mcu.c b/mt7915/mcu.c
index 4a709f9c..f2763247 100644
--- a/mt7915/mcu.c
+++ b/mt7915/mcu.c
@@ -2800,14 +2800,9 @@ int mt7915_mcu_set_chan_info(struct mt7915_phy *phy, int cmd)
 
 #ifdef CONFIG_NL80211_TESTMODE
 	if (phy->mt76->test.tx_antenna_mask &&
-	    (phy->mt76->test.state == MT76_TM_STATE_TX_FRAMES ||
-	     phy->mt76->test.state == MT76_TM_STATE_RX_FRAMES ||
-	     phy->mt76->test.state == MT76_TM_STATE_TX_CONT)) {
+	    mt76_testmode_enabled(phy->mt76)) {
 		req.tx_streams_num = fls(phy->mt76->test.tx_antenna_mask);
 		req.rx_streams = phy->mt76->test.tx_antenna_mask;
-
-		if (phy != &dev->phy)
-			req.rx_streams >>= dev->chainshift;
 	}
 #endif
 
diff --git a/mt7915/testmode.c b/mt7915/testmode.c
index 5e1767ab..a0360073 100644
--- a/mt7915/testmode.c
+++ b/mt7915/testmode.c
@@ -471,11 +471,7 @@ mt7915_tm_set_tx_frames(struct mt7915_phy *phy, bool en)
 		if (td->tx_spe_idx) {
 			phy->test.spe_idx = td->tx_spe_idx;
 		} else {
-			u8 tx_ant = td->tx_antenna_mask;
-
-			if (phy != &dev->phy)
-				tx_ant >>= dev->chainshift;
-			phy->test.spe_idx = spe_idx_map[tx_ant];
+			phy->test.spe_idx = spe_idx_map[td->tx_antenna_mask];
 		}
 	}
 
@@ -724,9 +720,6 @@ mt7915_tm_set_params(struct mt76_phy *mphy, struct nlattr **tb,
 	    td->state == MT76_TM_STATE_OFF)
 		return 0;
 
-	if (td->tx_antenna_mask & ~mphy->chainmask)
-		return -EINVAL;
-
 	for (i = 0; i < ARRAY_SIZE(tm_change_map); i++) {
 		if (tb[tm_change_map[i]])
 			changed |= BIT(i);
diff --git a/testmode.c b/testmode.c
index 382b4563..7cd00794 100644
--- a/testmode.c
+++ b/testmode.c
@@ -446,8 +446,8 @@ int mt76_testmode_cmd(struct ieee80211_hw *hw, struct ieee80211_vif *vif,
 	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_RATE_LDPC], &td->tx_rate_ldpc, 0, 1) ||
 	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_RATE_STBC], &td->tx_rate_stbc, 0, 1) ||
 	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_LTF], &td->tx_ltf, 0, 2) ||
-	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_ANTENNA],
-			   &td->tx_antenna_mask, 0, 0xff) ||
+	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_ANTENNA], &td->tx_antenna_mask,
+			   1, phy->antenna_mask) ||
 	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_SPE_IDX], &td->tx_spe_idx, 0, 27) ||
 	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_DUTY_CYCLE],
 			   &td->tx_duty_cycle, 0, 99) ||
-- 
2.18.0