Adding imu hostname to team_number parsing

Change-Id: I59caa6352e7de5c397c4b6d36123974c2104387b
Signed-off-by: Jim Ostrowski <yimmy13@gmail.com>
diff --git a/aos/network/team_number.cc b/aos/network/team_number.cc
index 175c3a9..bd90b96 100644
--- a/aos/network/team_number.cc
+++ b/aos/network/team_number.cc
@@ -130,6 +130,8 @@
     return std::string_view("pi");
   } else if (hostname.substr(0, 5) == "orin-") {
     return std::string_view("orin");
+  } else if (hostname.substr(0, 4) == "imu-") {
+    return std::string_view("orin");
   } else
     return std::nullopt;
 }
diff --git a/aos/network/team_number_test.cc b/aos/network/team_number_test.cc
index 9922b1b..278a364 100644
--- a/aos/network/team_number_test.cc
+++ b/aos/network/team_number_test.cc
@@ -28,7 +28,7 @@
   EXPECT_FALSE(ParseRoborioTeamNumber("roboRIO--FRC"));
 }
 
-TEST(TeamNumberTest, ParsePiOrOrinTeamNumber) {
+TEST(HostnameParseTest, ParsePiOrOrinTeamNumber) {
   EXPECT_EQ(971u, *ParsePiOrOrinTeamNumber("pi-971-1"));
   EXPECT_EQ(8971u, *ParsePiOrOrinTeamNumber("pi-8971-22"));
   EXPECT_EQ(8971u, *ParsePiOrOrinTeamNumber("pi-8971-"));
@@ -37,12 +37,16 @@
   EXPECT_EQ(8971u, *ParsePiOrOrinTeamNumber("orin-8971-22"));
   EXPECT_EQ(8971u, *ParsePiOrOrinTeamNumber("orin-8971-"));
 
+  EXPECT_FALSE(ParsePiOrOrinTeamNumber("roboRIO-971-FRC"));
+
   EXPECT_FALSE(ParseRoborioTeamNumber("pi"));
   EXPECT_FALSE(ParseRoborioTeamNumber("pi-"));
   EXPECT_FALSE(ParseRoborioTeamNumber("pi-971"));
   EXPECT_FALSE(ParseRoborioTeamNumber("pi-971a-1"));
   EXPECT_FALSE(ParseRoborioTeamNumber("orin-971-1"));
+}
 
+TEST(HostnameParseTest, ParsePiOrOrinNumber) {
   EXPECT_EQ(1u, *ParsePiOrOrinNumber("pi-971-1"));
   EXPECT_EQ(22u, *ParsePiOrOrinNumber("pi-8971-22"));
   EXPECT_EQ(1u, *ParsePiOrOrinNumber("orin-971-1"));
@@ -59,4 +63,19 @@
   EXPECT_FALSE(ParsePiOrOrinNumber("orin-971"));
 }
 
+TEST(HostnameParseTest, ParsePiOrOrin) {
+  EXPECT_EQ("pi", *ParsePiOrOrin("pi-971-1"));
+  EXPECT_EQ("pi", *ParsePiOrOrin("pi-8971-22"));
+  EXPECT_EQ("pi", *ParsePiOrOrin("pi-8971-"));
+
+  EXPECT_EQ("orin", *ParsePiOrOrin("orin-971-1"));
+  EXPECT_EQ("orin", *ParsePiOrOrin("orin-8971-22"));
+  EXPECT_EQ("orin", *ParsePiOrOrin("orin-8971-"));
+
+  EXPECT_EQ("orin", *ParsePiOrOrin("imu-971-1"));
+
+  EXPECT_FALSE(ParsePiOrOrin("roboRIO-971-FRC"));
+  EXPECT_FALSE(ParsePiOrOrin("laptop"));
+}
+
 }  // namespace aos::network::testing