improve debug flag checks for pytest

1 files changed, 3 insertions, 5 deletions

diff --git a/tests/test_discover.py b/tests/test_discover.py
index 22e1d73..c6d6ab2 100644
--- a/tests/test_discover.py
+++ b/tests/test_discover.py
@@ -1,6 +1,6 @@
 from unittest import IsolatedAsyncioTestCase
 from unittest.mock import AsyncMock, patch
-from discover import discover
+from discover import main
 from bleak.backends.device import BLEDevice
 import os
 
@@ -8,11 +8,10 @@ class TestDiscover(IsolatedAsyncioTestCase):
     @patch('src.discover.BleakScanner.discover', new_callable=AsyncMock)
     async def test_device_discovery_successful(self, mock_discover):
         device_1 = BLEDevice(address="AA:BB:CC:DD:EE:FF", name="TP350S", details={}, rssi=-40)
-        device_2 = BLEDevice(address="GG:HH:II:JJ:KK:LL", name="Device 2", details={}, rssi=-64)
 
         mock_discover.return_value = [device_1]
 
-        result = await discover()
+        result = await main(debug_flag=False)
 
         expected = [
             {"name": "TP350S", "address": "AA:BB:CC:DD:EE:FF"},
@@ -27,8 +26,7 @@ class TestDiscover(IsolatedAsyncioTestCase):
 
         mock_discover.return_value = [device_1, device_2]
 
-        with patch.dict(os.environ, {"DEBUG": "TRUE"}):
-            result = await discover()
+        result = await main(debug_flag=True)
 
         expected = [
             {"name": "TP350S", "address": "AA:BB:CC:DD:EE:FF"},