don't report `use_enums` on augmentations

Addresses part of https://github.com/dart-lang/linter/issues/4900

Change-Id: I9fd292ed29d5bddd9cafe1c3c376941183fcab76
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/355282
Reviewed-by: Konstantin Shcheglov <scheglov@google.com>
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Commit-Queue: Phil Quitslund <pquitslund@google.com>
diff --git a/pkg/linter/lib/src/extensions.dart b/pkg/linter/lib/src/extensions.dart
index 572281e..a0e4233 100644
--- a/pkg/linter/lib/src/extensions.dart
+++ b/pkg/linter/lib/src/extensions.dart
@@ -28,10 +28,12 @@
   bool get isAugmentation {
     var self = this;
     return switch (self) {
+      ClassDeclaration() => self.augmentKeyword != null,
       ConstructorDeclaration() => self.augmentKeyword != null,
       FunctionDeclarationImpl() => self.augmentKeyword != null,
       FunctionExpression() => self.parent?.isAugmentation ?? false,
       MethodDeclaration() => self.augmentKeyword != null,
+      MixinDeclaration() => self.augmentKeyword != null,
       // TODO(pq): unimplemented
       // VariableDeclaration() => self.augmentKeyword != null,
       _ => false
diff --git a/pkg/linter/lib/src/rules/use_enums.dart b/pkg/linter/lib/src/rules/use_enums.dart
index 74d0c80..bb74b27 100644
--- a/pkg/linter/lib/src/rules/use_enums.dart
+++ b/pkg/linter/lib/src/rules/use_enums.dart
@@ -9,6 +9,7 @@
 
 import '../analyzer.dart';
 import '../ast.dart';
+import '../extensions.dart';
 
 const _desc = r'Use enums rather than classes that behave like enums.';
 
@@ -190,6 +191,9 @@
 
   @override
   visitClassDeclaration(ClassDeclaration node) {
+    // Don't lint augmentations.
+    if (node.isAugmentation) return;
+
     if (node.abstractKeyword != null) return;
     var classElement = node.declaredElement;
     if (classElement == null) return;
diff --git a/pkg/linter/test/rules/use_enums_test.dart b/pkg/linter/test/rules/use_enums_test.dart
index be908d9..d78ab01 100644
--- a/pkg/linter/test/rules/use_enums_test.dart
+++ b/pkg/linter/test/rules/use_enums_test.dart
@@ -17,6 +17,25 @@
   @override
   String get lintRule => 'use_enums';
 
+  test_augmentation() async {
+    newFile('$testPackageLibPath/a.dart', r'''
+import augment 'test.dart';
+
+class C {}
+''');
+
+    await assertNoDiagnostics(r'''
+library augment 'a.dart';
+
+augment class C {
+  static const a = C._(1);
+  static const b = C._(2);
+  final int i;
+  const C._(this.i);
+}
+''');
+  }
+
   test_constructor_private() async {
     await assertDiagnostics(r'''
 class A {