fix: updating subscribe and publish to take same args for event topic

icp/python/src/triton_distributed/icp/nats_event_plane.py

@@ -246,7 +246,7 @@ class NatsEventPlane:
         self,
         payload: bytes | Any,
         event_type: Optional[str] = None,
-        event_topic: Optional[EventTopic] = None,
+        event_topic: Optional[EventTopic | str | List[str]] = None,
         timestamp: Optional[datetime.datetime] = datetime.datetime.now(datetime.UTC),
         event_id: Optional[uuid.UUID] = uuid.uuid4(),
     ) -> Event:
@@ -271,6 +271,9 @@ class NatsEventPlane:
         if event_id is None:
             event_id = uuid.uuid4()
 
+        if event_topic is not None and not isinstance(event_topic, EventTopic):
+            event_topic = EventTopic(event_topic)
+
         event_metadata = EventMetadata(
             event_id=event_id,
             event_topic=event_topic,
@@ -290,7 +293,6 @@ class NatsEventPlane:
 
         subject = self._compose_publish_subject(event_metadata)
         await self._nc.publish(subject, message)
-
         event_with_metadata = OnDemandEvent(
             payload, metadata_serialized, event_metadata
         )

icp/tests/python/event_plane/test_publish_subscribe.py

@@ -114,6 +114,32 @@ class TestEventPlaneFunctional:
         assert len(received_events) == 1
         assert received_events[0].event_id == event_metadata.event_id
 
+    @pytest.mark.asyncio
+    async def test_event_topic_list(self, nats_server, event_plane):
+        print(f"Print loop test: {id(asyncio.get_running_loop())}")
+
+        received_events: List[Event] = []
+
+        event = b"test_payload"
+
+        subscription = await event_plane.subscribe(event_topic="hello")
+        event_metadata = await event_plane.publish(event, event_topic=["hello"])
+
+        # Allow time for message to propagate
+        await asyncio.sleep(2)
+
+        async for x in subscription:
+            print(x.timestamp)
+            print(x.event_id)
+            print(x.event_type)
+            print(x.event_topic)
+            print(x.payload)
+            received_events.append(x)
+            break
+
+        assert len(received_events) == 1
+        assert received_events[0].event_id == event_metadata.event_id
+
     @pytest.mark.asyncio
     async def test_custom_type(self, nats_server, event_plane):
         print(f"Print loop test: {id(asyncio.get_running_loop())}")