Handle invalid country codes

Author: coreymckrill

includes/Generator/Customer.php

@@ -17,7 +17,7 @@ class Customer extends Generator {
 	 * @param bool  $save       Save the object before returning or not.
 	 * @param array $assoc_args Arguments passed via the CLI for additional customization.
 	 *
-	 * @return \WC_Customer Customer object with data populated.
+	 * @return \WC_Customer|\WP_Error Customer object with data populated.
 	 */
 	public static function generate( $save = true, array $assoc_args = array() ) {
 		self::init_faker();
@@ -43,6 +43,11 @@ public static function generate( $save = true, array $assoc_args = array() ) {
 
 		list( 'country' => $country, 'type' => $type ) = $args;
 
+		$country = CustomerInfo::get_valid_country_code( $country );
+		if ( is_wp_error( $country ) ) {
+			return $country;
+		}
+
 		if ( ! $type ) {
 			$type = self::$faker->randomDigit() < 7 ? 'person' : 'company'; // 70% person, 30% company.
 		}
@@ -140,6 +145,9 @@ public static function batch( $amount, array $args = array() ) {
 
 		for ( $i = 1; $i <= $amount; $i++ ) {
 			$customer       = self::generate( true, $args );
+			if ( is_wp_error( $customer ) ) {
+				return $customer;
+			}
 			$customer_ids[] = $customer->get_id();
 		}
 

includes/Generator/CustomerInfo.php

@@ -15,7 +15,7 @@ class CustomerInfo {
 	 *
 	 * @return string|\WP_Error
 	 */
-	protected static function get_valid_country_code( string $country_code = '' ) {
+	public static function get_valid_country_code( string $country_code = '' ) {
 		$country_code = strtoupper( $country_code );
 
 		if ( $country_code && ! WC()->countries->country_exists( $country_code ) ) {