Addressed PR comments

TBD54566975 · Aug 8, 2024 · 5358454 · 5358454
1 parent ff61129
commit 5358454
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/packages/protocol/src/message-kinds/order.ts b/packages/protocol/src/message-kinds/order.ts
@@ -16,7 +16,7 @@ export type CreateOrderOptions = {
  */
 export class Order extends Message {
   /** a set of valid Message kinds that can come after an order */
-  readonly validNext = new Set<MessageKind>(['orderinstructions'])
+  readonly validNext = new Set<MessageKind>(['orderinstructions', 'close'])
 
   /** The message kind (order) */
   readonly kind = 'order'

diff --git a/packages/protocol/tests/exchange.spec.ts b/packages/protocol/tests/exchange.spec.ts
@@ -273,20 +273,27 @@ describe('Exchange', () => {
         }
       })
 
-      it('can add an OrderInstructions after Order', async () => {
+      it('can add an OrderInstructions or a Close after Order', async () => {
         const exchange = new Exchange()
-
         exchange.addMessages([rfq, quote, order])
 
         exchange.addNextMessage(orderInstructions)
         expect(exchange.orderInstructions).to.deep.eq(orderInstructions)
       })
 
-      it('cannot add Rfq, Quote, Order, OrderStatus, or Close after Order', async () => {
+      it('can add a Close after Order', async () => {
+        const exchange = new Exchange()
+        exchange.addMessages([rfq, quote, order])
+
+        exchange.addNextMessage(closeByPfi)
+        expect(exchange.close).to.deep.eq(closeByPfi)
+      })
+
+      it('cannot add Rfq, Quote, Order, or OrderStatus after Order', async () => {
         const exchange = new Exchange()
         exchange.addMessages([rfq, quote, order])
 
-        for (const message of [rfq, quote, order, orderStatus, closeByAlice]) {
+        for (const message of [rfq, quote, order, orderStatus]) {
           try {
             exchange.addNextMessage(message)
             expect.fail()