Improve wp menu location assign command

Author: Peter J. Herrel

features/menu-location.feature

@@ -29,3 +29,23 @@ Feature: Manage WordPress menu locations
     Then STDOUT should be a table containing rows:
       | slug            | locations       |
       | primary-menu    |                 |
+
+    When I try `wp menu location assign secondary-menu secondary`
+    Then STDERR should be:
+      """
+      Error: Invalid menu secondary-menu.
+      """
+
+    When I run `wp menu create "Secondary Menu"`
+    And I try `wp menu location assign secondary-menu secondary`
+    Then STDERR should be:
+      """
+      Error: Invalid location secondary.
+      """
+
+    When I run `wp menu location assign secondary-menu primary`
+    And I run `wp menu list --fields=slug,locations`
+    Then STDOUT should be a table containing rows:
+      | slug            | locations       |
+      | primary-menu    |                 |
+      | secondary-menu  | primary         |

php/commands/menu-location.php

@@ -108,22 +108,22 @@ public function assign( $args, $_ ) {
 
 		list( $menu, $location ) = $args;
 
-		$menu = wp_get_nav_menu_object( $menu );
-		if ( ! $menu || is_wp_error( $menu ) ) {
-			WP_CLI::error( "Invalid menu." );
+		$menu_obj = wp_get_nav_menu_object( $menu );
+		if ( ! $menu_obj ) {
+			WP_CLI::error( "Invalid menu $menu." );
 		}
 
 		$locations = get_registered_nav_menus();
 		if ( ! array_key_exists( $location, $locations ) ) {
-			WP_CLI::error( "Invalid location." );
+			WP_CLI::error( "Invalid location $location." );
 		}
 
 		$locations = get_nav_menu_locations();
-		$locations[ $location ] = $menu->term_id;
+		$locations[ $location ] = $menu_obj->term_id;
 
 		set_theme_mod( 'nav_menu_locations', $locations );
 
-		WP_CLI::success( "Assigned location to menu." );
+		WP_CLI::success( "Assigned location $location to menu $menu." );
 	}
 
 	/**