From c893771bda3a0cad6153e70e09a36eedcfe8402f Mon Sep 17 00:00:00 2001
From: Andreas Lind <andreaslindpetersen@gmail.com>
Date: Sun, 16 Jun 2019 09:00:20 +0200
Subject: [PATCH] Take compound assertions into account when replacing 'to
 equal snapshot' with 'to inspect as snapshot'

---
 lib/ensureAfterHookIsRegistered.js |  28 +++++---
 test/unexpected-snapshot.js        | 110 +++++++++++++++++++++++------
 2 files changed, 109 insertions(+), 29 deletions(-)

diff --git a/lib/ensureAfterHookIsRegistered.js b/lib/ensureAfterHookIsRegistered.js
index 4ff77df..4db8405 100644
--- a/lib/ensureAfterHookIsRegistered.js
+++ b/lib/ensureAfterHookIsRegistered.js
@@ -228,19 +228,29 @@ function ensureAfterBlockIsRegistered(topLevelFixes) {
                 inspect
               ).replace(/\n^(?=[^\n])/gm, `\n${indent}`);
             }
+            // Take compound assertions into account:
+            const assertionArgument =
+              node.arguments[
+                node.arguments.length - (status === 'missing' ? 1 : 2)
+              ];
+            const newAssertionNameWithPrefix = assertionArgument.value.replace(
+              assertionName,
+              newAssertionName
+            );
+
             if (status === 'missing') {
-              if (newAssertionName !== assertionName) {
+              if (newAssertionName === assertionName) {
                 fixes.unshift(
-                  ruleFixer.replaceText(
-                    node.arguments[node.arguments.length - 1],
-                    `'${newAssertionName}', ${stringifiedSubject}`
+                  ruleFixer.insertTextAfter(
+                    assertionArgument,
+                    `, ${stringifiedSubject}`
                   )
                 );
               } else {
                 fixes.unshift(
-                  ruleFixer.insertTextAfter(
-                    node.arguments[node.arguments.length - 1],
-                    `, ${stringifiedSubject}`
+                  ruleFixer.replaceText(
+                    assertionArgument,
+                    `'${newAssertionNameWithPrefix}', ${stringifiedSubject}`
                   )
                 );
               }
@@ -248,8 +258,8 @@ function ensureAfterBlockIsRegistered(topLevelFixes) {
               if (newAssertionName !== assertionName) {
                 fixes.unshift(
                   ruleFixer.replaceText(
-                    node.arguments[node.arguments.length - 1],
-                    `'${newAssertionName}'`
+                    assertionArgument,
+                    `'${newAssertionNameWithPrefix}'`
                   )
                 );
               }
diff --git a/test/unexpected-snapshot.js b/test/unexpected-snapshot.js
index 5841d2f..96b5e3c 100644
--- a/test/unexpected-snapshot.js
+++ b/test/unexpected-snapshot.js
@@ -487,28 +487,98 @@ it('should foo', function() {
       );
     });
 
-    it('should switch to "to inspect as snapshot" when the subject contains a circular reference', function() {
-      return expect(
-        () => {
-          it('should foo', function() {
-            const foo = { bar: 123 };
-            foo.quux = foo;
-            expect(foo, 'to equal snapshot');
-          });
-        },
-        'to come out as',
-        function() {
-          it('should foo', function() {
-            const foo = { bar: 123 };
-            foo.quux = foo;
-            expect(
-              foo,
-              'to inspect as snapshot',
-              '{ bar: 123, quux: [Circular] }'
+    describe('when the subject contains a circular reference', function() {
+      it('should switch to "to inspect as snapshot"', function() {
+        return expect(
+          () => {
+            it('should foo', function() {
+              const foo = { bar: 123 };
+              foo.quux = foo;
+              expect(foo, 'to equal snapshot');
+            });
+          },
+          'to come out as',
+          function() {
+            it('should foo', function() {
+              const foo = { bar: 123 };
+              foo.quux = foo;
+              expect(
+                foo,
+                'to inspect as snapshot',
+                '{ bar: 123, quux: [Circular] }'
+              );
+            });
+          }
+        );
+      });
+
+      describe('with a compound assertion', function() {
+        it('should switch to "to inspect as snapshot"', function() {
+          return expect(
+            () => {
+              expect.addAssertion('<any> noop <assertion>', expect =>
+                expect.shift()
+              );
+
+              it('should foo', function() {
+                const foo = { bar: 123 };
+                foo.quux = foo;
+                expect(foo, 'noop to equal snapshot', { bar: 456 });
+              });
+            },
+            'to come out as',
+            function() {
+              expect.addAssertion('<any> noop <assertion>', expect =>
+                expect.shift()
+              );
+
+              it('should foo', function() {
+                const foo = { bar: 123 };
+                foo.quux = foo;
+                expect(
+                  foo,
+                  'noop to inspect as snapshot',
+                  '{ bar: 123, quux: [Circular] }'
+                );
+              });
+            }
+          );
+        });
+
+        describe('and no current snapshot', function() {
+          it('should switch to "to inspect as snapshot"', function() {
+            return expect(
+              () => {
+                expect.addAssertion('<any> noop <assertion>', expect =>
+                  expect.shift()
+                );
+
+                it('should foo', function() {
+                  const foo = { bar: 123 };
+                  foo.quux = foo;
+                  expect(foo, 'noop to equal snapshot');
+                });
+              },
+              'to come out as',
+              function() {
+                expect.addAssertion('<any> noop <assertion>', expect =>
+                  expect.shift()
+                );
+
+                it('should foo', function() {
+                  const foo = { bar: 123 };
+                  foo.quux = foo;
+                  expect(
+                    foo,
+                    'noop to inspect as snapshot',
+                    '{ bar: 123, quux: [Circular] }'
+                  );
+                });
+              }
             );
           });
-        }
-      );
+        });
+      });
     });
 
     it('should switch to "to inspect as snapshot" when the subject contains an object that does not have Object.prototype as its constructor', function() {