From db76fd53aebaa297dd35c27da146592eb1ac29e1 Mon Sep 17 00:00:00 2001
From: Jeroen van der Heijden <jeroen@cesbit.com>
Date: Tue, 30 May 2023 11:49:19 +0200
Subject: [PATCH] Change behavior, add collector when not assigned on announce

---
 pylibagent/agent.py   | 23 ++++++++++++++++++++++-
 pylibagent/version.py |  2 +-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/pylibagent/agent.py b/pylibagent/agent.py
index 12c368b..a3f29c3 100644
--- a/pylibagent/agent.py
+++ b/pylibagent/agent.py
@@ -96,7 +96,7 @@ async def announce(self, asset_name: Optional[str] = None,
             async with ClientSession(headers=self._headers) as session:
                 async with session.get(
                         url,
-                        params={'field': 'name'},
+                        params={'fields': 'name', 'collectors': 'key'},
                         ssl=self.verify_ssl) as r:
                     if r.status != 200:
                         msg = await r.text()
@@ -104,6 +104,27 @@ async def announce(self, asset_name: Optional[str] = None,
 
                     resp = await r.json()
                     name = resp["name"]
+                    collectors = resp["collectors"]
+
+            for collector in collectors:
+                if collector['key'] == self.key:
+                    break
+            else:
+                # The collector is not assigned yet
+                url = _join(
+                    self.api_uri,
+                    f'asset/{self.asset_id}/collector/{self.key}')
+                try:
+                    async with ClientSession(headers=self._headers) as session:
+                        async with session.post(url, ssl=self.verify_ssl) as r:
+                            if r.status != 204:
+                                msg = await r.text()
+                                raise Exception(
+                                    f'{msg} (error code: {r.status})')
+                except Exception as e:
+                    msg = str(e) or type(e).__name__
+                    logging.error(f'failed to assign collector: {msg}')
+
             logging.info(f'announced agent {name} (Id: {self.asset_id})')
             return
 
diff --git a/pylibagent/version.py b/pylibagent/version.py
index 1c98a23..7fd229a 100644
--- a/pylibagent/version.py
+++ b/pylibagent/version.py
@@ -1 +1 @@
-__version__ = '0.1.9'
+__version__ = '0.2.0'