Merge pull request #4 from RobertD502/query_fix

Get cats query fix
RobertD502 · Dec 27, 2023 · 3209675 · 3209675
2 parents ddae690 + 2899f36
commit 3209675
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 29 deletions.
diff --git a/lavviebot/constants.py b/lavviebot/constants.py
@@ -25,21 +25,47 @@
 DISCOVER_LB = '''query PurrsongTabLocations {getLocations {...LocationInfo getIots
                  {id lavviebot {nickname __typename} lavvieTag {id nickname __typename}
                  lavvieScanner {nickname __typename} pet {id cat {catMainPhoto nickname __typename} __typename} __typename} __typename}}
-                 fragment LocationInfo on Location {id nickname locationRole __typename}'''
+                 fragment LocationInfo on Location {id nickname locationRole hasUnknownCat __typename}'''
 
 """ Query to discover cats """
-DISCOVER_CATS = "query PetMain($data: PurrsongPetMainArgs!) " \
-                "{purrsongPetMain(data: $data) {location" \
-                "{id nickname hasUnknownCat hasLavviebot hasLavvieBox locationRole otherLocations" \
-                "{...LocationInfo __typename} __typename} selectedPet {...SelectedPet hasLavvieTag petClips" \
-                "{isDetectedOutlier __typename} __typename } otherPets {...OtherPetsForPetMain __typename } conditionScore dailyTotal" \
-                "{poopCount rest grooming walk run woodadaCount __typename} hourlyData" \
-                "{poopCount rest grooming walk run woodadaCount mainData __typename} __typename}} fragment LocationInfo on Location" \
-                "{id nickname locationRole __typename} fragment SelectedPet on Pet {id petType petCode cat" \
-                "{id nickname catSex catMainPhoto catAge catBirthDate catBirthDateCertainty __typename} dog" \
-                "{id nickname dogSex dogMainPhoto dogAge __typename} __typename} fragment OtherPetsForPetMain on Pet" \
-                "{id petType hasLavvieTag cat {id nickname catSex catMainPhoto catAge __typename} dog" \
-                "{id nickname dogSex dogMainPhoto __typename} petClips {isDetectedOutlier __typename} __typename}"
+
+DISCOVER_CATS = """
+query CatMain($locationId: Int, $includeLavvieCare: Boolean = true, $includeLavvieTag: Boolean = true, $includeDetailCatInfo: Boolean = true, $includeLocation: Boolean = true) {
+  getPets(data: {locationId: $locationId}) {
+    ...PetList
+    __typename
+  }
+}
+fragment PetList on Pet {
+  id
+  petCode
+  lavvieTag @include(if: $includeLavvieTag) {
+    id
+    lavvieTagUid
+    iotId
+    __typename
+  }
+  cat {
+    id
+    nickname
+    catMainPhoto
+    lavvieCare @include(if: $includeLavvieCare) {
+      recentStartDate
+      status
+      __typename
+    }
+    catAge @include(if: $includeDetailCatInfo)
+    catSex @include(if: $includeDetailCatInfo)
+    catLifeStage @include(if: $includeDetailCatInfo)
+    catBirthDate @include(if: $includeDetailCatInfo)
+    catBirthDateCertainty @include(if: $includeDetailCatInfo)
+    __typename
+  }
+  locationId @include(if: $includeLocation)
+  __typename
+}
+"""
+
 
 """ Query to get status of specific litter box """
 LB_STATUS = "query GetLavviebotDetails($data: IotIdArgs!) {getIotDetail(data: $data) " \

diff --git a/lavviebot/lavviebot_client.py b/lavviebot/lavviebot_client.py
@@ -127,11 +127,13 @@ async def async_discover_cats(self, location_id: int) -> ClientResponse:
             'User-Agent': USER_AGENT
         }
         dc_payload = {
-            "operationName": "PetMain",
+            "operationName": "CatMain",
             "variables": {
-                "data": {
-                    "locationId": location_id
-                }
+                "includeLavvieCare": True,
+                "includeLavvieTag": True,
+                "includeDetailCatInfo": True,
+                "includeLocation": True,
+                "locationId": location_id
             },
             "query": DISCOVER_CATS
         }
@@ -280,21 +282,16 @@ async def async_get_data(self) -> LavviebotData:
         if self.has_cat:
             for location in locations:
                 response = await self.async_discover_cats(location['id'])
-                location_data = response['data']['purrsongPetMain']['location']
-                if location_data['hasUnknownCat']:
+                if location['hasUnknownCat']:
                     unknown_cat = {
-                        'id': location_data['id'],
+                        'id': location['id'],
                         'is_unknown': True
                     }
                     cats.append(unknown_cat)
-                """ Also append the default cat. """
-                if selected_cat := response['data']['purrsongPetMain']['selectedPet']:
-                    selected_cat['is_unknown'] = False
-                    cats.append(selected_cat)
-
-                for other_cat in response['data']['purrsongPetMain']['otherPets']:
-                    other_cat['is_unknown'] = False
-                    cats.append(other_cat)
+                """ Append all cats to cat list. """
+                for cat in response['data']['getPets']:
+                    cat["is_unknown"] = False
+                    cats.append(cat)
 
             cat: dict
             for cat in cats:
@@ -351,6 +348,7 @@ async def async_get_data(self) -> LavviebotData:
         return LavviebotData(litterboxes=litter_box_data, cats=cat_data)
 
 
+
     async def async_fetch_all_endpoints(self, device_id: int) -> tuple[
         BaseException | Any, BaseException | Any, BaseException | Any]:
         """
@@ -628,3 +626,4 @@ async def _response(resp: ClientResponse, is_cookie: bool) -> SimpleCookie | Cli
         except Exception as e:
             raise LavviebotError(f'Could not return json: {e}') from e
         return response
+
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="lavviebotaio",
-    version="0.1",
+    version="0.2.0",
     author="Robert Drinovac",
     author_email="[email protected]",
     description="Asynchronous Python library for the PurrSong API utilized by LavvieBot S litter boxes",