#define SANDBOX_PCI_CLASS_CODE         PCI_CLASS_CODE_COMM
 #define SANDBOX_PCI_CLASS_SUB_CODE     PCI_CLASS_SUB_CODE_COMM_SERIAL
 
+/**
+ * sandbox_i2c_set_test_mode() - set test mode for running unit tests
+ *
+ * See sandbox_i2c_xfer() for the behaviour changes.
+ *
+ * @bus:       sandbox I2C bus to adjust
+ * @test_mode: true to select test mode, false to run normally
+ */
+void sandbox_i2c_set_test_mode(struct udevice *bus, bool test_mode);
+
 enum sandbox_i2c_eeprom_test_mode {
        SIE_TEST_MODE_NONE,
        /* Permits read/write of only one byte per I2C transaction */
 
 
 DECLARE_GLOBAL_DATA_PTR;
 
-struct dm_sandbox_i2c_emul_priv {
-       struct udevice *emul;
+struct sandbox_i2c_priv {
+       bool test_mode;
 };
 
 static int get_emul(struct udevice *dev, struct udevice **devp,
        return 0;
 }
 
+void sandbox_i2c_set_test_mode(struct udevice *bus, bool test_mode)
+{
+       struct sandbox_i2c_priv *priv = dev_get_priv(bus);
+
+       priv->test_mode = test_mode;
+}
+
 static int sandbox_i2c_xfer(struct udevice *bus, struct i2c_msg *msg,
                            int nmsgs)
 {
        struct dm_i2c_bus *i2c = dev_get_uclass_priv(bus);
+       struct sandbox_i2c_priv *priv = dev_get_priv(bus);
        struct dm_i2c_ops *ops;
        struct udevice *emul, *dev;
        bool is_read;
        int ret;
 
        /* Special test code to return success but with no emulation */
-       if (msg->addr == SANDBOX_I2C_TEST_ADDR)
+       if (priv->test_mode && msg->addr == SANDBOX_I2C_TEST_ADDR)
                return 0;
 
        ret = i2c_get_chip(bus, msg->addr, 1, &dev);
        if (ret)
                return ret;
 
-       /*
-        * For testing, don't allow writing above 100KHz for writes and
-        * 400KHz for reads
-        */
-       is_read = nmsgs > 1;
-       if (i2c->speed_hz > (is_read ? 400000 : 100000)) {
-               debug("%s: Max speed exceeded\n", __func__);
-               return -EINVAL;
+       if (priv->test_mode) {
+               /*
+               * For testing, don't allow writing above 100KHz for writes and
+               * 400KHz for reads.
+               */
+               is_read = nmsgs > 1;
+               if (i2c->speed_hz > (is_read ? 400000 : 100000)) {
+                       debug("%s: Max speed exceeded\n", __func__);
+                       return -EINVAL;
+               }
        }
+
        return ops->xfer(emul, msg, nmsgs);
 }
 
        .id     = UCLASS_I2C,
        .of_match = sandbox_i2c_ids,
        .ops    = &sandbox_i2c_ops,
+       .priv_auto_alloc_size = sizeof(struct sandbox_i2c_priv),
 };
 
        uint flags;
 #ifdef CONFIG_SANDBOX
        struct udevice *emul;
+       bool test_mode;
 #endif
 };
 
 
        uint8_t buf[5];
 
        ut_assertok(uclass_get_device_by_seq(UCLASS_I2C, busnum, &bus));
+
+       /* Use test mode so we create the required errors for invalid speeds */
+       sandbox_i2c_set_test_mode(bus, true);
        ut_assertok(i2c_get_chip(bus, chip, 1, &dev));
        ut_assertok(dm_i2c_set_bus_speed(bus, 100000));
        ut_assertok(dm_i2c_read(dev, 0, buf, 5));
        ut_asserteq(400000, dm_i2c_get_bus_speed(bus));
        ut_assertok(dm_i2c_read(dev, 0, buf, 5));
        ut_asserteq(-EINVAL, dm_i2c_write(dev, 0, buf, 5));
+       sandbox_i2c_set_test_mode(bus, false);
 
        return 0;
 }
        struct udevice *bus, *dev;
 
        ut_assertok(uclass_get_device_by_seq(UCLASS_I2C, busnum, &bus));
+
+       /* Use test mode so that this chip address will always probe */
+       sandbox_i2c_set_test_mode(bus, true);
        ut_assertok(dm_i2c_probe(bus, SANDBOX_I2C_TEST_ADDR, 0, &dev));
+       sandbox_i2c_set_test_mode(bus, false);
 
        return 0;
 }