Skip to content

Commit dcd438d

Browse files
j-gaocopybara-github
j-gao
authored and
copybara-github
committed
Internal change
PiperOrigin-RevId: 893307482
1 parent 4571f17 commit dcd438d

9 files changed

Lines changed: 226 additions & 0 deletions

File tree

‎src/java/com/google/devtools/mobileharness/infra/master/central/storage/lab/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ java_library(
4343
"DeviceRepository.java",
4444
],
4545
deps = [
46+
"//src/devtools/mobileharness/infra/master/central/proto:device_java_proto",
4647
"//src/java/com/google/devtools/mobileharness/api/model/error",
4748
"//src/java/com/google/devtools/mobileharness/api/model/lab:locator",
4849
"//src/java/com/google/devtools/mobileharness/shared/storage/transaction:transaction_context",

‎src/java/com/google/devtools/mobileharness/infra/master/central/storage/lab/DeviceRepository.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,29 @@
1717
package com.google.devtools.mobileharness.infra.master.central.storage.lab;
1818

1919
import com.google.devtools.mobileharness.api.model.error.MobileHarnessException;
20+
import com.google.devtools.mobileharness.api.model.lab.DeviceLocator;
2021
import com.google.devtools.mobileharness.api.model.lab.LabLocator;
22+
import com.google.devtools.mobileharness.infra.master.central.proto.Device.DeviceCondition;
2123
import com.google.devtools.mobileharness.shared.storage.transaction.TransactionContext;
24+
import java.util.Optional;
2225

2326
/** Interface for device repository. */
2427
public interface DeviceRepository {
2528

2629
/** Checks whether there is any device in the lab. */
2730
boolean hasDevice(LabLocator labLocator, TransactionContext context)
2831
throws MobileHarnessException;
32+
33+
/** Gets the condition of a device. */
34+
Optional<DeviceCondition> getDeviceCondition(
35+
DeviceLocator deviceLocator, TransactionContext context) throws MobileHarnessException;
36+
37+
/** Updates the condition of a device. */
38+
void updateDeviceCondition(
39+
DeviceLocator deviceLocator, DeviceCondition condition, TransactionContext context)
40+
throws MobileHarnessException;
41+
42+
/** Removes a device. */
43+
void removeDevice(DeviceLocator deviceLocator, TransactionContext context)
44+
throws MobileHarnessException;
2945
}

‎src/java/com/google/devtools/mobileharness/infra/master/central/storage/mysql/BUILD‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,9 +81,13 @@ java_library(
8181
name = "my_sql_device_table_client",
8282
srcs = ["MySqlDeviceTableClient.java"],
8383
deps = [
84+
"//src/devtools/mobileharness/infra/master/central/proto:device_java_proto",
8485
"//src/java/com/google/devtools/mobileharness/api/model/error",
8586
"//src/java/com/google/devtools/mobileharness/api/model/lab:locator",
87+
"//src/java/com/google/devtools/mobileharness/shared/util/base:proto_extension_registry",
88+
"//src/java/com/google/devtools/mobileharness/shared/util/logging:google_logger",
8689
"@maven//:javax_inject_jsr330_api",
90+
"@protobuf//:protobuf_java",
8791
],
8892
)
8993

@@ -92,6 +96,7 @@ java_library(
9296
srcs = ["MySqlDeviceRepository.java"],
9397
deps = [
9498
":my_sql_device_table_client",
99+
"//src/devtools/mobileharness/infra/master/central/proto:device_java_proto",
95100
"//src/java/com/google/devtools/mobileharness/api/model/error",
96101
"//src/java/com/google/devtools/mobileharness/api/model/lab:locator",
97102
"//src/java/com/google/devtools/mobileharness/infra/master/central/storage/lab:device_repository",

‎src/java/com/google/devtools/mobileharness/infra/master/central/storage/mysql/MySqlDeviceRepository.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,13 @@
1717
package com.google.devtools.mobileharness.infra.master.central.storage.mysql;
1818

1919
import com.google.devtools.mobileharness.api.model.error.MobileHarnessException;
20+
import com.google.devtools.mobileharness.api.model.lab.DeviceLocator;
2021
import com.google.devtools.mobileharness.api.model.lab.LabLocator;
22+
import com.google.devtools.mobileharness.infra.master.central.proto.Device.DeviceCondition;
2123
import com.google.devtools.mobileharness.infra.master.central.storage.lab.DeviceRepository;
2224
import com.google.devtools.mobileharness.shared.storage.transaction.TransactionContext;
2325
import com.google.devtools.mobileharness.shared.storage.transaction.mysql.MySqlTransactionContext;
26+
import java.util.Optional;
2427
import javax.inject.Inject;
2528

2629
/** MySQL implementation of {@link DeviceRepository}. */
@@ -40,6 +43,29 @@ public boolean hasDevice(LabLocator labLocator, TransactionContext transactionCo
4043
return mySqlDeviceTableClient.hasDevice(labLocator, context.getConnection());
4144
}
4245

46+
@Override
47+
public Optional<DeviceCondition> getDeviceCondition(
48+
DeviceLocator deviceLocator, TransactionContext transactionContext)
49+
throws MobileHarnessException {
50+
MySqlTransactionContext context = castContext(transactionContext);
51+
return mySqlDeviceTableClient.getDeviceCondition(deviceLocator, context.getConnection());
52+
}
53+
54+
@Override
55+
public void updateDeviceCondition(
56+
DeviceLocator deviceLocator, DeviceCondition condition, TransactionContext transactionContext)
57+
throws MobileHarnessException {
58+
MySqlTransactionContext context = castContext(transactionContext);
59+
mySqlDeviceTableClient.updateDeviceCondition(deviceLocator, condition, context.getConnection());
60+
}
61+
62+
@Override
63+
public void removeDevice(DeviceLocator deviceLocator, TransactionContext transactionContext)
64+
throws MobileHarnessException {
65+
MySqlTransactionContext context = castContext(transactionContext);
66+
mySqlDeviceTableClient.removeDevice(deviceLocator, context.getConnection());
67+
}
68+
4369
private MySqlTransactionContext castContext(TransactionContext context) {
4470
if (!(context instanceof MySqlTransactionContext mysqlTransactionContext)) {
4571
throw new IllegalArgumentException(

‎src/java/com/google/devtools/mobileharness/infra/master/central/storage/mysql/MySqlDeviceTableClient.java‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,19 +16,29 @@
1616

1717
package com.google.devtools.mobileharness.infra.master.central.storage.mysql;
1818

19+
import com.google.common.flogger.FluentLogger;
1920
import com.google.devtools.mobileharness.api.model.error.BasicErrorId;
2021
import com.google.devtools.mobileharness.api.model.error.MobileHarnessException;
22+
import com.google.devtools.mobileharness.api.model.lab.DeviceLocator;
2123
import com.google.devtools.mobileharness.api.model.lab.LabLocator;
24+
import com.google.devtools.mobileharness.infra.master.central.proto.Device.DeviceCondition;
25+
import com.google.devtools.mobileharness.shared.util.base.ProtoExtensionRegistry;
26+
import com.google.protobuf.InvalidProtocolBufferException;
2227
import java.sql.Connection;
2328
import java.sql.PreparedStatement;
2429
import java.sql.ResultSet;
2530
import java.sql.SQLException;
31+
import java.util.Optional;
2632
import javax.inject.Inject;
2733

2834
/** Client for reading/writing device table in MySQL. */
2935
public class MySqlDeviceTableClient {
36+
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
37+
3038
private static final String TABLE_NAME = "DeviceInfo";
3139
private static final String COL_LAB_ID = "LabId";
40+
private static final String COL_DEVICE_ID = "DeviceId";
41+
private static final String COL_DEVICE_CONDITION = "DeviceCondition";
3242

3343
@Inject
3444
MySqlDeviceTableClient() {}
@@ -49,4 +59,76 @@ public boolean hasDevice(LabLocator labLocator, Connection connection)
4959
e);
5060
}
5161
}
62+
63+
/** Gets the condition of a device. Empty if the device or column value is null. */
64+
public Optional<DeviceCondition> getDeviceCondition(
65+
DeviceLocator deviceLocator, Connection connection) throws MobileHarnessException {
66+
String sql =
67+
String.format(
68+
"SELECT %s FROM %s WHERE %s = ? AND %s = ?",
69+
COL_DEVICE_CONDITION, TABLE_NAME, COL_LAB_ID, COL_DEVICE_ID);
70+
try (PreparedStatement preparedStatement = connection.prepareStatement(sql)) {
71+
preparedStatement.setString(1, deviceLocator.labLocator().getId());
72+
preparedStatement.setString(2, deviceLocator.id());
73+
try (ResultSet resultSet = preparedStatement.executeQuery()) {
74+
if (resultSet.next()) {
75+
byte[] conditionBytes = resultSet.getBytes(COL_DEVICE_CONDITION);
76+
if (conditionBytes == null) {
77+
return Optional.empty();
78+
}
79+
return Optional.of(
80+
DeviceCondition.parseFrom(
81+
conditionBytes, ProtoExtensionRegistry.getGeneratedRegistry()));
82+
} else {
83+
return Optional.empty();
84+
}
85+
}
86+
} catch (SQLException | InvalidProtocolBufferException e) {
87+
throw new MobileHarnessException(
88+
BasicErrorId.DATABASE_TABLE_READ_ERROR,
89+
"Failed to get DeviceCondition for " + deviceLocator,
90+
e);
91+
}
92+
}
93+
94+
/** Updates the condition of a device. */
95+
public void updateDeviceCondition(
96+
DeviceLocator deviceLocator, DeviceCondition condition, Connection connection)
97+
throws MobileHarnessException {
98+
String sql =
99+
String.format(
100+
"UPDATE %s SET %s = ? WHERE %s = ? AND %s = ?",
101+
TABLE_NAME, COL_DEVICE_CONDITION, COL_LAB_ID, COL_DEVICE_ID);
102+
try (PreparedStatement preparedStatement = connection.prepareStatement(sql)) {
103+
preparedStatement.setBytes(1, condition.toByteArray());
104+
preparedStatement.setString(2, deviceLocator.labLocator().getId());
105+
preparedStatement.setString(3, deviceLocator.id());
106+
int rowsAffected = preparedStatement.executeUpdate();
107+
if (rowsAffected == 0) {
108+
logger.atWarning().log(
109+
"DeviceCondition not updated, DeviceId %s not found.", deviceLocator.id());
110+
}
111+
} catch (SQLException e) {
112+
throw new MobileHarnessException(
113+
BasicErrorId.DATABASE_TABLE_UPDATE_ERROR,
114+
"Failed to update DeviceCondition for " + deviceLocator,
115+
e);
116+
}
117+
}
118+
119+
/** Removes a device. */
120+
public void removeDevice(DeviceLocator deviceLocator, Connection connection)
121+
throws MobileHarnessException {
122+
String sql =
123+
String.format(
124+
"DELETE FROM %s WHERE %s = ? AND %s = ?", TABLE_NAME, COL_LAB_ID, COL_DEVICE_ID);
125+
try (PreparedStatement preparedStatement = connection.prepareStatement(sql)) {
126+
preparedStatement.setString(1, deviceLocator.labLocator().getId());
127+
preparedStatement.setString(2, deviceLocator.id());
128+
preparedStatement.executeUpdate();
129+
} catch (SQLException e) {
130+
throw new MobileHarnessException(
131+
BasicErrorId.DATABASE_TABLE_DELETE_ERROR, "Failed to delete device " + deviceLocator, e);
132+
}
133+
}
52134
}

‎src/java/com/google/devtools/mobileharness/shared/storage/transaction/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ java_library(
4141
":transaction_config",
4242
":transaction_worker",
4343
"//src/java/com/google/devtools/mobileharness/api/model/error",
44+
"@maven//:com_google_errorprone_error_prone_annotations",
4445
],
4546
)
4647

‎src/java/com/google/devtools/mobileharness/shared/storage/transaction/TransactionRunner.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package com.google.devtools.mobileharness.shared.storage.transaction;
1818

1919
import com.google.devtools.mobileharness.api.model.error.MobileHarnessException;
20+
import com.google.errorprone.annotations.CanIgnoreReturnValue;
2021

2122
/** Runner for transactions. */
2223
public interface TransactionRunner {
@@ -29,6 +30,7 @@ public interface TransactionRunner {
2930
* @return the result of the transaction
3031
* @throws MobileHarnessException if fails to run the transaction
3132
*/
33+
@CanIgnoreReturnValue
3234
public abstract <R> R run(TransactionConfig config, TransactionWorker<R> worker)
3335
throws MobileHarnessException;
3436
}

‎src/javatests/com/google/devtools/mobileharness/infra/master/central/storage/mysql/BUILD‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ java_library(
2525
name = "mysql",
2626
srcs = glob(["*.java"]),
2727
deps = [
28+
"//src/devtools/mobileharness/api/model/proto:device_java_proto",
29+
"//src/devtools/mobileharness/infra/master/central/proto:device_java_proto",
2830
"//src/devtools/mobileharness/infra/master/central/proto:lab_java_proto",
2931
"//src/java/com/google/devtools/mobileharness/api/model/error",
3032
"//src/java/com/google/devtools/mobileharness/api/model/lab:locator",

‎src/javatests/com/google/devtools/mobileharness/infra/master/central/storage/mysql/MySqlDeviceTableClientTest.java‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,15 @@
2323
import static org.mockito.Mockito.when;
2424

2525
import com.google.devtools.mobileharness.api.model.error.MobileHarnessException;
26+
import com.google.devtools.mobileharness.api.model.lab.DeviceLocator;
2627
import com.google.devtools.mobileharness.api.model.lab.LabLocator;
28+
import com.google.devtools.mobileharness.api.model.proto.Device.DeviceStatus;
29+
import com.google.devtools.mobileharness.infra.master.central.proto.Device.DeviceCondition;
2730
import java.sql.Connection;
2831
import java.sql.PreparedStatement;
2932
import java.sql.ResultSet;
3033
import java.sql.SQLException;
34+
import java.util.Optional;
3135
import org.junit.Before;
3236
import org.junit.Rule;
3337
import org.junit.Test;
@@ -49,6 +53,10 @@ public final class MySqlDeviceTableClientTest {
4953
private static final String LAB_ID = "test_lab_id";
5054
private static final LabLocator LAB_LOCATOR =
5155
LabLocator.of(/* ip= */ "test_lab_ip", /* hostName= */ LAB_ID);
56+
private static final String DEVICE_ID = "test_device_id";
57+
private static final DeviceLocator DEVICE_LOCATOR = DeviceLocator.of(DEVICE_ID, LAB_LOCATOR);
58+
private static final DeviceCondition DEVICE_CONDITION =
59+
DeviceCondition.newBuilder().setStatusFromLab(DeviceStatus.MISSING).build();
5260

5361
private MySqlDeviceTableClient mySqlDeviceTableClient;
5462

@@ -90,4 +98,87 @@ public void hasDevice_sqlException() throws Exception {
9098
() -> mySqlDeviceTableClient.hasDevice(LAB_LOCATOR, mockConnection));
9199
assertThat(exception).hasMessageThat().contains("Failed to check devices for lab");
92100
}
101+
102+
@Test
103+
public void getDeviceCondition_success() throws Exception {
104+
when(mockResultSet.next()).thenReturn(true);
105+
when(mockResultSet.getBytes("DeviceCondition")).thenReturn(DEVICE_CONDITION.toByteArray());
106+
107+
Optional<DeviceCondition> result =
108+
mySqlDeviceTableClient.getDeviceCondition(DEVICE_LOCATOR, mockConnection);
109+
110+
assertThat(result).isPresent();
111+
assertThat(result.get()).isEqualTo(DEVICE_CONDITION);
112+
verify(mockPreparedStatement).setString(1, LAB_ID);
113+
verify(mockPreparedStatement).setString(2, DEVICE_ID);
114+
}
115+
116+
@Test
117+
public void getDeviceCondition_empty() throws Exception {
118+
when(mockResultSet.next()).thenReturn(false);
119+
120+
Optional<DeviceCondition> result =
121+
mySqlDeviceTableClient.getDeviceCondition(DEVICE_LOCATOR, mockConnection);
122+
123+
assertThat(result).isEmpty();
124+
}
125+
126+
@Test
127+
public void getDeviceCondition_nullValue() throws Exception {
128+
when(mockResultSet.next()).thenReturn(true);
129+
when(mockResultSet.getBytes("DeviceCondition")).thenReturn(null);
130+
131+
Optional<DeviceCondition> result =
132+
mySqlDeviceTableClient.getDeviceCondition(DEVICE_LOCATOR, mockConnection);
133+
134+
assertThat(result).isEmpty();
135+
}
136+
137+
@Test
138+
public void getDeviceCondition_exception() throws Exception {
139+
when(mockPreparedStatement.executeQuery()).thenThrow(new SQLException("DB error"));
140+
141+
assertThrows(
142+
MobileHarnessException.class,
143+
() -> mySqlDeviceTableClient.getDeviceCondition(DEVICE_LOCATOR, mockConnection));
144+
}
145+
146+
@Test
147+
public void updateDeviceCondition_success() throws Exception {
148+
mySqlDeviceTableClient.updateDeviceCondition(DEVICE_LOCATOR, DEVICE_CONDITION, mockConnection);
149+
150+
verify(mockPreparedStatement).setBytes(1, DEVICE_CONDITION.toByteArray());
151+
verify(mockPreparedStatement).setString(2, LAB_ID);
152+
verify(mockPreparedStatement).setString(3, DEVICE_ID);
153+
verify(mockPreparedStatement).executeUpdate();
154+
}
155+
156+
@Test
157+
public void updateDeviceCondition_exception() throws Exception {
158+
when(mockPreparedStatement.executeUpdate()).thenThrow(new SQLException("DB error"));
159+
160+
assertThrows(
161+
MobileHarnessException.class,
162+
() ->
163+
mySqlDeviceTableClient.updateDeviceCondition(
164+
DEVICE_LOCATOR, DEVICE_CONDITION, mockConnection));
165+
}
166+
167+
@Test
168+
public void removeDevice_success() throws Exception {
169+
mySqlDeviceTableClient.removeDevice(DEVICE_LOCATOR, mockConnection);
170+
171+
verify(mockPreparedStatement).setString(1, LAB_ID);
172+
verify(mockPreparedStatement).setString(2, DEVICE_ID);
173+
verify(mockPreparedStatement).executeUpdate();
174+
}
175+
176+
@Test
177+
public void removeDevice_exception() throws Exception {
178+
when(mockPreparedStatement.executeUpdate()).thenThrow(new SQLException("DB error"));
179+
180+
assertThrows(
181+
MobileHarnessException.class,
182+
() -> mySqlDeviceTableClient.removeDevice(DEVICE_LOCATOR, mockConnection));
183+
}
93184
}

0 commit comments

Comments
 (0)